Compare commits
11 Commits
main
...
33d11ae223
| Author | SHA1 | Date | |
|---|---|---|---|
| 33d11ae223 | |||
| 05e90d3e2b | |||
| fe414d49e6 | |||
| d002d35bde | |||
| c9c3d17db0 | |||
| a909455f97 | |||
| 67381b02db | |||
| 235f84fa19 | |||
| 9c777c8429 | |||
| 0b17a0f4c8 | |||
| 2eabe55fe6 |
@@ -11,9 +11,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
- Comprehensive documentation suite including guides for architecture, configuration, testing, and more.
|
- Comprehensive documentation suite including guides for architecture, configuration, testing, and more.
|
||||||
- Rustdoc examples for core components like `Provider` and `SessionController`.
|
- Rustdoc examples for core components like `Provider` and `SessionController`.
|
||||||
- Module-level documentation for `owlen-tui`.
|
- Module-level documentation for `owlen-tui`.
|
||||||
|
- Ollama integration can now talk to Ollama Cloud when an API key is configured.
|
||||||
|
- Ollama provider will also read `OLLAMA_API_KEY` / `OLLAMA_CLOUD_API_KEY` environment variables when no key is stored in the config.
|
||||||
|
|
||||||
### Changed
|
### Changed
|
||||||
- The main `README.md` has been updated to be more concise and link to the new documentation.
|
- The main `README.md` has been updated to be more concise and link to the new documentation.
|
||||||
|
- Default configuration now pre-populates both `providers.ollama` and `providers.ollama-cloud` entries so switching between local and cloud backends is a single setting change.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
18
Cargo.toml
18
Cargo.toml
@@ -5,6 +5,9 @@ members = [
|
|||||||
"crates/owlen-tui",
|
"crates/owlen-tui",
|
||||||
"crates/owlen-cli",
|
"crates/owlen-cli",
|
||||||
"crates/owlen-ollama",
|
"crates/owlen-ollama",
|
||||||
|
"crates/owlen-mcp-server",
|
||||||
|
"crates/owlen-mcp-llm-server",
|
||||||
|
"crates/owlen-mcp-client",
|
||||||
]
|
]
|
||||||
exclude = []
|
exclude = []
|
||||||
|
|
||||||
@@ -34,12 +37,24 @@ tui-textarea = "0.6"
|
|||||||
# HTTP client and JSON handling
|
# HTTP client and JSON handling
|
||||||
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] }
|
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] }
|
||||||
serde = { version = "1.0", features = ["derive"] }
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
serde_json = "1.0"
|
serde_json = { version = "1.0" }
|
||||||
|
|
||||||
# Utilities
|
# Utilities
|
||||||
uuid = { version = "1.0", features = ["v4", "serde"] }
|
uuid = { version = "1.0", features = ["v4", "serde"] }
|
||||||
anyhow = "1.0"
|
anyhow = "1.0"
|
||||||
thiserror = "1.0"
|
thiserror = "1.0"
|
||||||
|
nix = "0.29"
|
||||||
|
which = "6.0"
|
||||||
|
tempfile = "3.8"
|
||||||
|
jsonschema = "0.17"
|
||||||
|
aes-gcm = "0.10"
|
||||||
|
ring = "0.17"
|
||||||
|
keyring = "3.0"
|
||||||
|
chrono = { version = "0.4", features = ["serde"] }
|
||||||
|
urlencoding = "2.1"
|
||||||
|
regex = "1.10"
|
||||||
|
rpassword = "7.3"
|
||||||
|
sqlx = { version = "0.7", default-features = false, features = ["runtime-tokio-rustls", "sqlite", "macros", "uuid", "chrono", "migrate"] }
|
||||||
|
|
||||||
# Configuration
|
# Configuration
|
||||||
toml = "0.8"
|
toml = "0.8"
|
||||||
@@ -58,7 +73,6 @@ async-trait = "0.1"
|
|||||||
clap = { version = "4.0", features = ["derive"] }
|
clap = { version = "4.0", features = ["derive"] }
|
||||||
|
|
||||||
# Dev dependencies
|
# Dev dependencies
|
||||||
tempfile = "3.8"
|
|
||||||
tokio-test = "0.4"
|
tokio-test = "0.4"
|
||||||
|
|
||||||
# For more keys and their definitions, see https://doc.rust-lang.org/cargo/reference/manifest.html
|
# For more keys and their definitions, see https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ description = "Command-line interface for OWLEN LLM client"
|
|||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["chat-client"]
|
default = ["chat-client"]
|
||||||
chat-client = []
|
chat-client = ["owlen-tui"]
|
||||||
code-client = []
|
code-client = []
|
||||||
|
|
||||||
[[bin]]
|
[[bin]]
|
||||||
@@ -23,10 +23,16 @@ name = "owlen-code"
|
|||||||
path = "src/code_main.rs"
|
path = "src/code_main.rs"
|
||||||
required-features = ["code-client"]
|
required-features = ["code-client"]
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "owlen-agent"
|
||||||
|
path = "src/agent_main.rs"
|
||||||
|
required-features = ["chat-client"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
owlen-core = { path = "../owlen-core" }
|
owlen-core = { path = "../owlen-core" }
|
||||||
owlen-tui = { path = "../owlen-tui" }
|
|
||||||
owlen-ollama = { path = "../owlen-ollama" }
|
owlen-ollama = { path = "../owlen-ollama" }
|
||||||
|
# Optional TUI dependency, enabled by the "chat-client" feature.
|
||||||
|
owlen-tui = { path = "../owlen-tui", optional = true }
|
||||||
|
|
||||||
# CLI framework
|
# CLI framework
|
||||||
clap = { version = "4.0", features = ["derive"] }
|
clap = { version = "4.0", features = ["derive"] }
|
||||||
@@ -43,3 +49,6 @@ crossterm = { workspace = true }
|
|||||||
anyhow = { workspace = true }
|
anyhow = { workspace = true }
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
serde_json = { workspace = true }
|
serde_json = { workspace = true }
|
||||||
|
regex = "1"
|
||||||
|
thiserror = "1"
|
||||||
|
dirs = "5"
|
||||||
|
|||||||
54
crates/owlen-cli/src/agent_main.rs
Normal file
54
crates/owlen-cli/src/agent_main.rs
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
//! Simple entry point for the ReAct agentic executor.
|
||||||
|
//!
|
||||||
|
//! Usage: `owlen-agent "<prompt>" [--model <model>] [--max-iter <n>]`
|
||||||
|
//!
|
||||||
|
//! This binary demonstrates Phase 4 without the full TUI. It creates an
|
||||||
|
//! OllamaProvider, a RemoteMcpClient, runs the AgentExecutor and prints the
|
||||||
|
//! final answer.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use clap::Parser;
|
||||||
|
use owlen_cli::agent::{AgentConfig, AgentExecutor};
|
||||||
|
use owlen_core::mcp::remote_client::RemoteMcpClient;
|
||||||
|
use owlen_ollama::OllamaProvider;
|
||||||
|
|
||||||
|
/// Command‑line arguments for the agent binary.
|
||||||
|
#[derive(Parser, Debug)]
|
||||||
|
#[command(name = "owlen-agent", author, version, about = "Run the ReAct agent")]
|
||||||
|
struct Args {
|
||||||
|
/// The initial user query.
|
||||||
|
prompt: String,
|
||||||
|
/// Model to use (defaults to Ollama default).
|
||||||
|
#[arg(long)]
|
||||||
|
model: Option<String>,
|
||||||
|
/// Maximum ReAct iterations.
|
||||||
|
#[arg(long, default_value_t = 10)]
|
||||||
|
max_iter: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
let args = Args::parse();
|
||||||
|
|
||||||
|
// Initialise the LLM provider (Ollama) – uses default local URL.
|
||||||
|
let provider = Arc::new(OllamaProvider::new("http://localhost:11434")?);
|
||||||
|
// Initialise the MCP client (remote LLM server) – this client also knows how
|
||||||
|
// to call the built‑in resource tools.
|
||||||
|
let mcp_client = Arc::new(RemoteMcpClient::new()?);
|
||||||
|
|
||||||
|
let config = AgentConfig {
|
||||||
|
max_iterations: args.max_iter,
|
||||||
|
model: args.model.unwrap_or_else(|| "llama3.2:latest".to_string()),
|
||||||
|
..AgentConfig::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let executor = AgentExecutor::new(provider, mcp_client, config, None);
|
||||||
|
match executor.run(args.prompt).await {
|
||||||
|
Ok(answer) => {
|
||||||
|
println!("\nFinal answer:\n{}", answer);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(e) => Err(anyhow::anyhow!(e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use clap::{Arg, Command};
|
use clap::{Arg, Command};
|
||||||
use owlen_core::session::SessionController;
|
use owlen_core::{session::SessionController, storage::StorageManager};
|
||||||
use owlen_ollama::OllamaProvider;
|
use owlen_ollama::OllamaProvider;
|
||||||
use owlen_tui::{config, ui, AppState, CodeApp, Event, EventHandler, SessionEvent};
|
use owlen_tui::{config, ui, AppState, CodeApp, Event, EventHandler, SessionEvent};
|
||||||
use std::io;
|
use std::io;
|
||||||
@@ -17,7 +17,7 @@ use crossterm::{
|
|||||||
};
|
};
|
||||||
use ratatui::{backend::CrosstermBackend, Terminal};
|
use ratatui::{backend::CrosstermBackend, Terminal};
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main(flavor = "multi_thread")]
|
||||||
async fn main() -> Result<()> {
|
async fn main() -> Result<()> {
|
||||||
let matches = Command::new("owlen-code")
|
let matches = Command::new("owlen-code")
|
||||||
.about("OWLEN Code Mode - TUI optimized for programming assistance")
|
.about("OWLEN Code Mode - TUI optimized for programming assistance")
|
||||||
@@ -32,19 +32,42 @@ async fn main() -> Result<()> {
|
|||||||
.get_matches();
|
.get_matches();
|
||||||
|
|
||||||
let mut config = config::try_load_config().unwrap_or_default();
|
let mut config = config::try_load_config().unwrap_or_default();
|
||||||
|
// Disable encryption for code mode.
|
||||||
|
config.privacy.encrypt_local_data = false;
|
||||||
|
|
||||||
if let Some(model) = matches.get_one::<String>("model") {
|
if let Some(model) = matches.get_one::<String>("model") {
|
||||||
config.general.default_model = Some(model.clone());
|
config.general.default_model = Some(model.clone());
|
||||||
}
|
}
|
||||||
|
|
||||||
let provider_cfg = config::ensure_ollama_config(&mut config).clone();
|
let provider_name = config.general.default_provider.clone();
|
||||||
|
let provider_cfg = config::ensure_provider_config(&mut config, &provider_name).clone();
|
||||||
|
|
||||||
|
let provider_type = provider_cfg.provider_type.to_ascii_lowercase();
|
||||||
|
if provider_type != "ollama" && provider_type != "ollama-cloud" {
|
||||||
|
anyhow::bail!(
|
||||||
|
"Unsupported provider type '{}' configured for provider '{}'",
|
||||||
|
provider_cfg.provider_type,
|
||||||
|
provider_name
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
let provider = Arc::new(OllamaProvider::from_config(
|
let provider = Arc::new(OllamaProvider::from_config(
|
||||||
&provider_cfg,
|
&provider_cfg,
|
||||||
Some(&config.general),
|
Some(&config.general),
|
||||||
)?);
|
)?);
|
||||||
|
|
||||||
let controller = SessionController::new(provider, config.clone());
|
let storage = Arc::new(StorageManager::new().await?);
|
||||||
let (mut app, mut session_rx) = CodeApp::new(controller);
|
// Code client - code execution tools enabled
|
||||||
|
use owlen_core::ui::NoOpUiController;
|
||||||
|
let controller = SessionController::new(
|
||||||
|
provider,
|
||||||
|
config.clone(),
|
||||||
|
storage.clone(),
|
||||||
|
Arc::new(NoOpUiController),
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
let (mut app, mut session_rx) = CodeApp::new(controller).await?;
|
||||||
app.inner_mut().initialize_models().await?;
|
app.inner_mut().initialize_models().await?;
|
||||||
|
|
||||||
let cancellation_token = CancellationToken::new();
|
let cancellation_token = CancellationToken::new();
|
||||||
@@ -63,7 +86,7 @@ async fn main() -> Result<()> {
|
|||||||
cancellation_token.cancel();
|
cancellation_token.cancel();
|
||||||
event_handle.await?;
|
event_handle.await?;
|
||||||
|
|
||||||
config::save_config(app.inner().config())?;
|
config::save_config(&app.inner().config())?;
|
||||||
|
|
||||||
disable_raw_mode()?;
|
disable_raw_mode()?;
|
||||||
execute!(
|
execute!(
|
||||||
@@ -87,8 +110,21 @@ async fn run_app(
|
|||||||
session_rx: &mut mpsc::UnboundedReceiver<SessionEvent>,
|
session_rx: &mut mpsc::UnboundedReceiver<SessionEvent>,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
loop {
|
loop {
|
||||||
|
// Advance loading animation frame
|
||||||
|
app.inner_mut().advance_loading_animation();
|
||||||
|
|
||||||
terminal.draw(|f| ui::render_chat(f, app.inner_mut()))?;
|
terminal.draw(|f| ui::render_chat(f, app.inner_mut()))?;
|
||||||
|
|
||||||
|
// Process any pending LLM requests AFTER UI has been drawn
|
||||||
|
if let Err(e) = app.inner_mut().process_pending_llm_request().await {
|
||||||
|
eprintln!("Error processing LLM request: {}", e);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Process any pending tool executions AFTER UI has been drawn
|
||||||
|
if let Err(e) = app.inner_mut().process_pending_tool_execution().await {
|
||||||
|
eprintln!("Error processing tool execution: {}", e);
|
||||||
|
}
|
||||||
|
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
Some(event) = event_rx.recv() => {
|
Some(event) = event_rx.recv() => {
|
||||||
if let AppState::Quit = app.handle_event(event).await? {
|
if let AppState::Quit = app.handle_event(event).await? {
|
||||||
@@ -98,6 +134,10 @@ async fn run_app(
|
|||||||
Some(session_event) = session_rx.recv() => {
|
Some(session_event) = session_rx.recv() => {
|
||||||
app.handle_session_event(session_event)?;
|
app.handle_session_event(session_event)?;
|
||||||
}
|
}
|
||||||
|
// Add a timeout to keep the animation going even when there are no events
|
||||||
|
_ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {
|
||||||
|
// This will cause the loop to continue and advance the animation
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
8
crates/owlen-cli/src/lib.rs
Normal file
8
crates/owlen-cli/src/lib.rs
Normal file
@@ -0,0 +1,8 @@
|
|||||||
|
//! Library portion of the `owlen-cli` crate.
|
||||||
|
//!
|
||||||
|
//! It currently only re‑exports the `agent` module used by the standalone
|
||||||
|
//! `owlen-agent` binary. Additional shared functionality can be added here in
|
||||||
|
//! the future.
|
||||||
|
|
||||||
|
// Re-export agent module from owlen-core
|
||||||
|
pub use owlen_core::agent;
|
||||||
@@ -1,9 +1,9 @@
|
|||||||
//! OWLEN CLI - Chat TUI client
|
//! OWLEN CLI - Chat TUI client
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use clap::{Arg, Command};
|
use owlen_core::{session::SessionController, storage::StorageManager};
|
||||||
use owlen_core::session::SessionController;
|
|
||||||
use owlen_ollama::OllamaProvider;
|
use owlen_ollama::OllamaProvider;
|
||||||
|
use owlen_tui::tui_controller::{TuiController, TuiRequest};
|
||||||
use owlen_tui::{config, ui, AppState, ChatApp, Event, EventHandler, SessionEvent};
|
use owlen_tui::{config, ui, AppState, ChatApp, Event, EventHandler, SessionEvent};
|
||||||
use std::io;
|
use std::io;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@@ -15,37 +15,42 @@ use crossterm::{
|
|||||||
execute,
|
execute,
|
||||||
terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
|
terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
|
||||||
};
|
};
|
||||||
use ratatui::{backend::CrosstermBackend, Terminal};
|
use ratatui::{prelude::CrosstermBackend, Terminal};
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main(flavor = "multi_thread")]
|
||||||
async fn main() -> Result<()> {
|
async fn main() -> Result<()> {
|
||||||
let matches = Command::new("owlen")
|
// (imports completed above)
|
||||||
.about("OWLEN - A chat-focused TUI client for Ollama")
|
|
||||||
.version(env!("CARGO_PKG_VERSION"))
|
|
||||||
.arg(
|
|
||||||
Arg::new("model")
|
|
||||||
.short('m')
|
|
||||||
.long("model")
|
|
||||||
.value_name("MODEL")
|
|
||||||
.help("Preferred model to use for this session"),
|
|
||||||
)
|
|
||||||
.get_matches();
|
|
||||||
|
|
||||||
let mut config = config::try_load_config().unwrap_or_default();
|
// (main logic starts below)
|
||||||
|
// Set auto-consent for TUI mode to prevent blocking stdin reads
|
||||||
|
std::env::set_var("OWLEN_AUTO_CONSENT", "1");
|
||||||
|
|
||||||
if let Some(model) = matches.get_one::<String>("model") {
|
let (tui_tx, _tui_rx) = mpsc::unbounded_channel::<TuiRequest>();
|
||||||
config.general.default_model = Some(model.clone());
|
let tui_controller = Arc::new(TuiController::new(tui_tx));
|
||||||
|
|
||||||
|
// Load configuration (or fall back to defaults) for the session controller.
|
||||||
|
let mut cfg = config::try_load_config().unwrap_or_default();
|
||||||
|
// Disable encryption for CLI to avoid password prompts in this environment.
|
||||||
|
cfg.privacy.encrypt_local_data = false;
|
||||||
|
// Determine provider configuration
|
||||||
|
let provider_name = cfg.general.default_provider.clone();
|
||||||
|
let provider_cfg = config::ensure_provider_config(&mut cfg, &provider_name).clone();
|
||||||
|
let provider_type = provider_cfg.provider_type.to_ascii_lowercase();
|
||||||
|
if provider_type != "ollama" && provider_type != "ollama-cloud" {
|
||||||
|
anyhow::bail!(
|
||||||
|
"Unsupported provider type '{}' configured for provider '{}'",
|
||||||
|
provider_cfg.provider_type,
|
||||||
|
provider_name,
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Prepare provider from configuration
|
|
||||||
let provider_cfg = config::ensure_ollama_config(&mut config).clone();
|
|
||||||
let provider = Arc::new(OllamaProvider::from_config(
|
let provider = Arc::new(OllamaProvider::from_config(
|
||||||
&provider_cfg,
|
&provider_cfg,
|
||||||
Some(&config.general),
|
Some(&cfg.general),
|
||||||
)?);
|
)?);
|
||||||
|
let storage = Arc::new(StorageManager::new().await?);
|
||||||
let controller = SessionController::new(provider, config.clone());
|
let controller =
|
||||||
let (mut app, mut session_rx) = ChatApp::new(controller);
|
SessionController::new(provider, cfg, storage.clone(), tui_controller, false).await?;
|
||||||
|
let (mut app, mut session_rx) = ChatApp::new(controller).await?;
|
||||||
app.initialize_models().await?;
|
app.initialize_models().await?;
|
||||||
|
|
||||||
// Event infrastructure
|
// Event infrastructure
|
||||||
@@ -73,7 +78,7 @@ async fn main() -> Result<()> {
|
|||||||
event_handle.await?;
|
event_handle.await?;
|
||||||
|
|
||||||
// Persist configuration updates (e.g., selected model)
|
// Persist configuration updates (e.g., selected model)
|
||||||
config::save_config(app.config())?;
|
config::save_config(&app.config())?;
|
||||||
|
|
||||||
disable_raw_mode()?;
|
disable_raw_mode()?;
|
||||||
execute!(
|
execute!(
|
||||||
@@ -104,7 +109,14 @@ async fn run_app(
|
|||||||
terminal.draw(|f| ui::render_chat(f, app))?;
|
terminal.draw(|f| ui::render_chat(f, app))?;
|
||||||
|
|
||||||
// Process any pending LLM requests AFTER UI has been drawn
|
// Process any pending LLM requests AFTER UI has been drawn
|
||||||
app.process_pending_llm_request().await?;
|
if let Err(e) = app.process_pending_llm_request().await {
|
||||||
|
eprintln!("Error processing LLM request: {}", e);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Process any pending tool executions AFTER UI has been drawn
|
||||||
|
if let Err(e) = app.process_pending_tool_execution().await {
|
||||||
|
eprintln!("Error processing tool execution: {}", e);
|
||||||
|
}
|
||||||
|
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
Some(event) = event_rx.recv() => {
|
Some(event) = event_rx.recv() => {
|
||||||
|
|||||||
271
crates/owlen-cli/tests/agent_tests.rs
Normal file
271
crates/owlen-cli/tests/agent_tests.rs
Normal file
@@ -0,0 +1,271 @@
|
|||||||
|
//! Integration tests for the ReAct agent loop functionality.
|
||||||
|
//!
|
||||||
|
//! These tests verify that the agent executor correctly:
|
||||||
|
//! - Parses ReAct formatted responses
|
||||||
|
//! - Executes tool calls
|
||||||
|
//! - Handles multi-step workflows
|
||||||
|
//! - Recovers from errors
|
||||||
|
//! - Respects iteration limits
|
||||||
|
|
||||||
|
use owlen_cli::agent::{AgentConfig, AgentExecutor, LlmResponse};
|
||||||
|
use owlen_core::mcp::remote_client::RemoteMcpClient;
|
||||||
|
use owlen_ollama::OllamaProvider;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_react_parsing_tool_call() {
|
||||||
|
let executor = create_test_executor();
|
||||||
|
|
||||||
|
// Test parsing a tool call with JSON arguments
|
||||||
|
let text = "THOUGHT: I should search for information\nACTION: web_search\nACTION_INPUT: {\"query\": \"rust async programming\"}\n";
|
||||||
|
|
||||||
|
let result = executor.parse_response(text);
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(LlmResponse::ToolCall {
|
||||||
|
thought,
|
||||||
|
tool_name,
|
||||||
|
arguments,
|
||||||
|
}) => {
|
||||||
|
assert_eq!(thought, "I should search for information");
|
||||||
|
assert_eq!(tool_name, "web_search");
|
||||||
|
assert_eq!(arguments["query"], "rust async programming");
|
||||||
|
}
|
||||||
|
other => panic!("Expected ToolCall, got: {:?}", other),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_react_parsing_final_answer() {
|
||||||
|
let executor = create_test_executor();
|
||||||
|
|
||||||
|
let text = "THOUGHT: I have enough information now\nACTION: final_answer\nACTION_INPUT: The answer is 42\n";
|
||||||
|
|
||||||
|
let result = executor.parse_response(text);
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(LlmResponse::FinalAnswer { thought, answer }) => {
|
||||||
|
assert_eq!(thought, "I have enough information now");
|
||||||
|
assert_eq!(answer, "The answer is 42");
|
||||||
|
}
|
||||||
|
other => panic!("Expected FinalAnswer, got: {:?}", other),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_react_parsing_with_multiline_thought() {
|
||||||
|
let executor = create_test_executor();
|
||||||
|
|
||||||
|
let text = "THOUGHT: This is a complex\nmulti-line thought\nACTION: list_files\nACTION_INPUT: {\"path\": \".\"}\n";
|
||||||
|
|
||||||
|
let result = executor.parse_response(text);
|
||||||
|
|
||||||
|
// The regex currently only captures until first newline
|
||||||
|
// This test documents current behavior
|
||||||
|
match result {
|
||||||
|
Ok(LlmResponse::ToolCall { thought, .. }) => {
|
||||||
|
// Regex pattern stops at first \n after THOUGHT:
|
||||||
|
assert!(thought.contains("This is a complex"));
|
||||||
|
}
|
||||||
|
other => panic!("Expected ToolCall, got: {:?}", other),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore] // Requires Ollama to be running
|
||||||
|
async fn test_agent_single_tool_scenario() {
|
||||||
|
// This test requires a running Ollama instance and MCP server
|
||||||
|
let provider = Arc::new(OllamaProvider::new("http://localhost:11434").unwrap());
|
||||||
|
let mcp_client = Arc::new(RemoteMcpClient::new().unwrap());
|
||||||
|
|
||||||
|
let config = AgentConfig {
|
||||||
|
max_iterations: 5,
|
||||||
|
model: "llama3.2".to_string(),
|
||||||
|
temperature: Some(0.7),
|
||||||
|
max_tokens: None,
|
||||||
|
max_tool_calls: 10,
|
||||||
|
};
|
||||||
|
|
||||||
|
let executor = AgentExecutor::new(provider, mcp_client, config, None);
|
||||||
|
|
||||||
|
// Simple query that should complete in one tool call
|
||||||
|
let result = executor
|
||||||
|
.run("List files in the current directory".to_string())
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(answer) => {
|
||||||
|
assert!(!answer.is_empty(), "Answer should not be empty");
|
||||||
|
println!("Agent answer: {}", answer);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
// It's okay if this fails due to LLM not following format
|
||||||
|
println!("Agent test skipped: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore] // Requires Ollama to be running
|
||||||
|
async fn test_agent_multi_step_workflow() {
|
||||||
|
// Test a query that requires multiple tool calls
|
||||||
|
let provider = Arc::new(OllamaProvider::new("http://localhost:11434").unwrap());
|
||||||
|
let mcp_client = Arc::new(RemoteMcpClient::new().unwrap());
|
||||||
|
|
||||||
|
let config = AgentConfig {
|
||||||
|
max_iterations: 10,
|
||||||
|
model: "llama3.2".to_string(),
|
||||||
|
temperature: Some(0.5), // Lower temperature for more consistent behavior
|
||||||
|
max_tokens: None,
|
||||||
|
max_tool_calls: 20,
|
||||||
|
};
|
||||||
|
|
||||||
|
let executor = AgentExecutor::new(provider, mcp_client, config, None);
|
||||||
|
|
||||||
|
// Query requiring multiple steps: list -> read -> analyze
|
||||||
|
let result = executor
|
||||||
|
.run("Find all Rust files and tell me which one contains 'Agent'".to_string())
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(answer) => {
|
||||||
|
assert!(!answer.is_empty());
|
||||||
|
println!("Multi-step answer: {}", answer);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
println!("Multi-step test skipped: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore] // Requires Ollama
|
||||||
|
async fn test_agent_iteration_limit() {
|
||||||
|
let provider = Arc::new(OllamaProvider::new("http://localhost:11434").unwrap());
|
||||||
|
let mcp_client = Arc::new(RemoteMcpClient::new().unwrap());
|
||||||
|
|
||||||
|
let config = AgentConfig {
|
||||||
|
max_iterations: 2, // Very low limit to test enforcement
|
||||||
|
model: "llama3.2".to_string(),
|
||||||
|
temperature: Some(0.7),
|
||||||
|
max_tokens: None,
|
||||||
|
max_tool_calls: 5,
|
||||||
|
};
|
||||||
|
|
||||||
|
let executor = AgentExecutor::new(provider, mcp_client, config, None);
|
||||||
|
|
||||||
|
// Complex query that would require many iterations
|
||||||
|
let result = executor
|
||||||
|
.run("Perform an exhaustive analysis of all files".to_string())
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Should hit the iteration limit (or parse error if LLM doesn't follow format)
|
||||||
|
match result {
|
||||||
|
Err(e) => {
|
||||||
|
let error_str = format!("{}", e);
|
||||||
|
// Accept either iteration limit error or parse error (LLM didn't follow ReAct format)
|
||||||
|
assert!(
|
||||||
|
error_str.contains("Maximum iterations")
|
||||||
|
|| error_str.contains("2")
|
||||||
|
|| error_str.contains("parse"),
|
||||||
|
"Expected iteration limit or parse error, got: {}",
|
||||||
|
error_str
|
||||||
|
);
|
||||||
|
println!("Test passed: agent stopped with error: {}", error_str);
|
||||||
|
}
|
||||||
|
Ok(_) => {
|
||||||
|
// It's possible the LLM completed within 2 iterations
|
||||||
|
println!("Agent completed within iteration limit");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore] // Requires Ollama
|
||||||
|
async fn test_agent_tool_budget_enforcement() {
|
||||||
|
let provider = Arc::new(OllamaProvider::new("http://localhost:11434").unwrap());
|
||||||
|
let mcp_client = Arc::new(RemoteMcpClient::new().unwrap());
|
||||||
|
|
||||||
|
let config = AgentConfig {
|
||||||
|
max_iterations: 20,
|
||||||
|
model: "llama3.2".to_string(),
|
||||||
|
temperature: Some(0.7),
|
||||||
|
max_tokens: None,
|
||||||
|
max_tool_calls: 3, // Very low tool call budget
|
||||||
|
};
|
||||||
|
|
||||||
|
let executor = AgentExecutor::new(provider, mcp_client, config, None);
|
||||||
|
|
||||||
|
// Query that would require many tool calls
|
||||||
|
let result = executor
|
||||||
|
.run("Read every file in the project and summarize them all".to_string())
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Should hit the tool call budget (or parse error if LLM doesn't follow format)
|
||||||
|
match result {
|
||||||
|
Err(e) => {
|
||||||
|
let error_str = format!("{}", e);
|
||||||
|
// Accept either budget error or parse error (LLM didn't follow ReAct format)
|
||||||
|
assert!(
|
||||||
|
error_str.contains("Maximum iterations")
|
||||||
|
|| error_str.contains("budget")
|
||||||
|
|| error_str.contains("parse"),
|
||||||
|
"Expected budget or parse error, got: {}",
|
||||||
|
error_str
|
||||||
|
);
|
||||||
|
println!("Test passed: agent stopped with error: {}", error_str);
|
||||||
|
}
|
||||||
|
Ok(_) => {
|
||||||
|
println!("Agent completed within tool budget");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Helper function to create a test executor
|
||||||
|
// For parsing tests, we don't need a real connection
|
||||||
|
fn create_test_executor() -> AgentExecutor {
|
||||||
|
// Create dummy instances - the parse_response method doesn't actually use them
|
||||||
|
let provider = Arc::new(OllamaProvider::new("http://localhost:11434").unwrap());
|
||||||
|
|
||||||
|
// For parsing tests, we can accept the error from RemoteMcpClient::new()
|
||||||
|
// since we're only testing parse_response which doesn't use the MCP client
|
||||||
|
let mcp_client = match RemoteMcpClient::new() {
|
||||||
|
Ok(client) => Arc::new(client),
|
||||||
|
Err(_) => {
|
||||||
|
// If MCP server binary doesn't exist, parsing tests can still run
|
||||||
|
// by using a dummy client that will never be called
|
||||||
|
// This is a workaround for unit tests that only need parse_response
|
||||||
|
panic!("MCP server binary not found - build the project first with: cargo build --all");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let config = AgentConfig::default();
|
||||||
|
AgentExecutor::new(provider, mcp_client, config, None)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_agent_config_defaults() {
|
||||||
|
let config = AgentConfig::default();
|
||||||
|
|
||||||
|
assert_eq!(config.max_iterations, 10);
|
||||||
|
assert_eq!(config.model, "ollama");
|
||||||
|
assert_eq!(config.temperature, Some(0.7));
|
||||||
|
assert_eq!(config.max_tool_calls, 20);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_agent_config_custom() {
|
||||||
|
let config = AgentConfig {
|
||||||
|
max_iterations: 15,
|
||||||
|
model: "custom-model".to_string(),
|
||||||
|
temperature: Some(0.5),
|
||||||
|
max_tokens: Some(2000),
|
||||||
|
max_tool_calls: 30,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(config.max_iterations, 15);
|
||||||
|
assert_eq!(config.model, "custom-model");
|
||||||
|
assert_eq!(config.temperature, Some(0.5));
|
||||||
|
assert_eq!(config.max_tokens, Some(2000));
|
||||||
|
assert_eq!(config.max_tool_calls, 30);
|
||||||
|
}
|
||||||
@@ -9,23 +9,40 @@ homepage.workspace = true
|
|||||||
description = "Core traits and types for OWLEN LLM client"
|
description = "Core traits and types for OWLEN LLM client"
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
anyhow = "1.0.75"
|
anyhow = { workspace = true }
|
||||||
log = "0.4.20"
|
log = "0.4.20"
|
||||||
serde = { version = "1.0.188", features = ["derive"] }
|
regex = { workspace = true }
|
||||||
serde_json = "1.0.105"
|
serde = { workspace = true }
|
||||||
thiserror = "1.0.48"
|
serde_json = { workspace = true }
|
||||||
tokio = { version = "1.32.0", features = ["full"] }
|
thiserror = { workspace = true }
|
||||||
|
tokio = { workspace = true }
|
||||||
unicode-segmentation = "1.11"
|
unicode-segmentation = "1.11"
|
||||||
unicode-width = "0.1"
|
unicode-width = "0.1"
|
||||||
uuid = { version = "1.4.1", features = ["v4", "serde"] }
|
uuid = { workspace = true }
|
||||||
textwrap = "0.16.0"
|
textwrap = { workspace = true }
|
||||||
futures = "0.3.28"
|
futures = { workspace = true }
|
||||||
async-trait = "0.1.73"
|
async-trait = { workspace = true }
|
||||||
toml = "0.8.0"
|
toml = { workspace = true }
|
||||||
shellexpand = "3.1.0"
|
shellexpand = { workspace = true }
|
||||||
dirs = "5.0"
|
dirs = "5.0"
|
||||||
ratatui = { workspace = true }
|
ratatui = { workspace = true }
|
||||||
|
tempfile = { workspace = true }
|
||||||
|
jsonschema = { workspace = true }
|
||||||
|
which = { workspace = true }
|
||||||
|
nix = { workspace = true }
|
||||||
|
aes-gcm = { workspace = true }
|
||||||
|
ring = { workspace = true }
|
||||||
|
keyring = { workspace = true }
|
||||||
|
chrono = { workspace = true }
|
||||||
|
crossterm = { workspace = true }
|
||||||
|
urlencoding = { workspace = true }
|
||||||
|
rpassword = { workspace = true }
|
||||||
|
sqlx = { workspace = true }
|
||||||
|
duckduckgo = "0.2.0"
|
||||||
|
reqwest = { workspace = true, features = ["default"] }
|
||||||
|
reqwest_011 = { version = "0.11", package = "reqwest" }
|
||||||
|
path-clean = "1.0"
|
||||||
|
tokio-stream = "0.1"
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tokio-test = { workspace = true }
|
tokio-test = { workspace = true }
|
||||||
tempfile = { workspace = true }
|
|
||||||
|
|||||||
12
crates/owlen-core/migrations/0001_create_conversations.sql
Normal file
12
crates/owlen-core/migrations/0001_create_conversations.sql
Normal file
@@ -0,0 +1,12 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS conversations (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
name TEXT,
|
||||||
|
description TEXT,
|
||||||
|
model TEXT NOT NULL,
|
||||||
|
message_count INTEGER NOT NULL,
|
||||||
|
created_at INTEGER NOT NULL,
|
||||||
|
updated_at INTEGER NOT NULL,
|
||||||
|
data TEXT NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_conversations_updated_at ON conversations(updated_at DESC);
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS secure_items (
|
||||||
|
key TEXT PRIMARY KEY,
|
||||||
|
nonce BLOB NOT NULL,
|
||||||
|
ciphertext BLOB NOT NULL,
|
||||||
|
created_at INTEGER NOT NULL,
|
||||||
|
updated_at INTEGER NOT NULL
|
||||||
|
);
|
||||||
377
crates/owlen-core/src/agent.rs
Normal file
377
crates/owlen-core/src/agent.rs
Normal file
@@ -0,0 +1,377 @@
|
|||||||
|
//! High‑level agentic executor implementing the ReAct pattern.
|
||||||
|
//!
|
||||||
|
//! The executor coordinates three responsibilities:
|
||||||
|
//! 1. Build a ReAct prompt from the conversation history and the list of
|
||||||
|
//! available MCP tools.
|
||||||
|
//! 2. Send the prompt to an LLM provider (any type implementing
|
||||||
|
//! `owlen_core::Provider`).
|
||||||
|
//! 3. Parse the LLM response, optionally invoke a tool via an MCP client,
|
||||||
|
//! and feed the observation back into the conversation.
|
||||||
|
//!
|
||||||
|
//! The implementation is intentionally minimal – it provides the core loop
|
||||||
|
//! required by Phase 4 of the roadmap. Integration with the TUI and additional
|
||||||
|
//! safety mechanisms can be added on top of this module.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use crate::ui::UiController;
|
||||||
|
|
||||||
|
use dirs;
|
||||||
|
use regex::Regex;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::fs::OpenOptions;
|
||||||
|
use std::io::Write;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
use tokio::signal;
|
||||||
|
|
||||||
|
use crate::mcp::client::McpClient;
|
||||||
|
use crate::mcp::{McpToolCall, McpToolDescriptor, McpToolResponse};
|
||||||
|
use crate::{
|
||||||
|
types::{ChatRequest, Message},
|
||||||
|
Error, Provider, Result as CoreResult,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// Configuration for the agent executor.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct AgentConfig {
|
||||||
|
/// Maximum number of ReAct iterations before the executor aborts.
|
||||||
|
pub max_iterations: usize,
|
||||||
|
/// Model name to use for the LLM provider.
|
||||||
|
pub model: String,
|
||||||
|
/// Optional temperature.
|
||||||
|
pub temperature: Option<f32>,
|
||||||
|
/// Optional max_tokens.
|
||||||
|
pub max_tokens: Option<u32>,
|
||||||
|
/// Maximum number of tool calls allowed per execution (budget).
|
||||||
|
pub max_tool_calls: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AgentConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
max_iterations: 10,
|
||||||
|
model: "ollama".into(),
|
||||||
|
temperature: Some(0.7),
|
||||||
|
max_tokens: None,
|
||||||
|
max_tool_calls: 20,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Enum representing the possible parsed LLM responses in ReAct format.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum LlmResponse {
|
||||||
|
/// A reasoning step without action.
|
||||||
|
Reasoning { thought: String },
|
||||||
|
/// The model wants to invoke a tool.
|
||||||
|
ToolCall {
|
||||||
|
thought: String,
|
||||||
|
tool_name: String,
|
||||||
|
arguments: serde_json::Value,
|
||||||
|
},
|
||||||
|
/// The model produced a final answer.
|
||||||
|
FinalAnswer { thought: String, answer: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Error type for the agent executor.
|
||||||
|
#[derive(thiserror::Error, Debug)]
|
||||||
|
pub enum AgentError {
|
||||||
|
#[error("LLM provider error: {0}")]
|
||||||
|
Provider(Error),
|
||||||
|
#[error("MCP client error: {0}")]
|
||||||
|
Mcp(Error),
|
||||||
|
#[error("Tool execution denied by user")]
|
||||||
|
ToolDenied,
|
||||||
|
#[error("Failed to parse LLM response")]
|
||||||
|
Parse,
|
||||||
|
#[error("Maximum iterations ({0}) reached without final answer")]
|
||||||
|
MaxIterationsReached(usize),
|
||||||
|
#[error("Agent execution cancelled by user (Ctrl+C)")]
|
||||||
|
Cancelled,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Core executor handling the ReAct loop.
|
||||||
|
pub struct AgentExecutor {
|
||||||
|
llm_client: Arc<dyn Provider + Send + Sync>,
|
||||||
|
tool_client: Arc<dyn McpClient + Send + Sync>,
|
||||||
|
config: AgentConfig,
|
||||||
|
ui_controller: Option<Arc<dyn UiController + Send + Sync>>, // optional UI for confirmations
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AgentExecutor {
|
||||||
|
/// Construct a new executor.
|
||||||
|
pub fn new(
|
||||||
|
llm_client: Arc<dyn Provider + Send + Sync>,
|
||||||
|
tool_client: Arc<dyn McpClient + Send + Sync>,
|
||||||
|
config: AgentConfig,
|
||||||
|
ui_controller: Option<Arc<dyn UiController + Send + Sync>>, // pass None for headless
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
llm_client,
|
||||||
|
tool_client,
|
||||||
|
config,
|
||||||
|
ui_controller,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Discover tools exposed by the MCP server.
|
||||||
|
async fn discover_tools(&self) -> CoreResult<Vec<McpToolDescriptor>> {
|
||||||
|
self.tool_client.list_tools().await
|
||||||
|
}
|
||||||
|
|
||||||
|
// #[allow(dead_code)]
|
||||||
|
// Build a ReAct prompt from the current message history and discovered tools.
|
||||||
|
/*
|
||||||
|
#[allow(dead_code)]
|
||||||
|
fn build_prompt(
|
||||||
|
&self,
|
||||||
|
history: &[Message],
|
||||||
|
tools: &[McpToolDescriptor],
|
||||||
|
) -> String {
|
||||||
|
// System prompt describing the format.
|
||||||
|
let system = "You are an intelligent agent following the ReAct pattern. Use the following sections:\nTHOUGHT: your reasoning\nACTION: the tool name you want to call (or "final_answer")\nACTION_INPUT: JSON arguments for the tool.\nIf ACTION is "final_answer", provide the final answer in the next line after the ACTION_INPUT.\n";
|
||||||
|
|
||||||
|
let mut prompt = format!("System: {}\n", system);
|
||||||
|
// Append conversation history.
|
||||||
|
for msg in history {
|
||||||
|
let role = match msg.role {
|
||||||
|
Role::User => "User",
|
||||||
|
Role::Assistant => "Assistant",
|
||||||
|
Role::System => "System",
|
||||||
|
Role::Tool => "Tool",
|
||||||
|
};
|
||||||
|
prompt.push_str(&format!("{}: {}\n", role, msg.content));
|
||||||
|
}
|
||||||
|
// Append tool descriptions.
|
||||||
|
if !tools.is_empty() {
|
||||||
|
let tools_json = json!(tools);
|
||||||
|
prompt.push_str(&format!("Available tools (JSON schema): {}\n", tools_json));
|
||||||
|
}
|
||||||
|
prompt
|
||||||
|
}
|
||||||
|
*/
|
||||||
|
|
||||||
|
// build_prompt removed; not used in current implementation
|
||||||
|
|
||||||
|
/// Parse raw LLM text into a structured `LlmResponse`.
|
||||||
|
pub fn parse_response(&self, text: &str) -> std::result::Result<LlmResponse, AgentError> {
|
||||||
|
// Normalise line endings.
|
||||||
|
let txt = text.trim();
|
||||||
|
// Regex patterns for parsing ReAct format.
|
||||||
|
// THOUGHT and ACTION capture up to the next newline.
|
||||||
|
// ACTION_INPUT captures everything remaining (including multiline JSON).
|
||||||
|
let thought_re = Regex::new(r"(?s)THOUGHT:\s*(?P<thought>.+?)(?:\n|$)").unwrap();
|
||||||
|
let action_re = Regex::new(r"(?s)ACTION:\s*(?P<action>.+?)(?:\n|$)").unwrap();
|
||||||
|
// ACTION_INPUT captures rest of text (multiline-friendly)
|
||||||
|
let input_re = Regex::new(r"(?s)ACTION_INPUT:\s*(?P<input>.+)").unwrap();
|
||||||
|
|
||||||
|
let thought = thought_re
|
||||||
|
.captures(txt)
|
||||||
|
.and_then(|c| c.name("thought"))
|
||||||
|
.map(|m| m.as_str().trim().to_string())
|
||||||
|
.ok_or(AgentError::Parse)?;
|
||||||
|
let action = action_re
|
||||||
|
.captures(txt)
|
||||||
|
.and_then(|c| c.name("action"))
|
||||||
|
.map(|m| m.as_str().trim().to_string())
|
||||||
|
.ok_or(AgentError::Parse)?;
|
||||||
|
let input = input_re
|
||||||
|
.captures(txt)
|
||||||
|
.and_then(|c| c.name("input"))
|
||||||
|
.map(|m| m.as_str().trim().to_string())
|
||||||
|
.ok_or(AgentError::Parse)?;
|
||||||
|
|
||||||
|
if action.eq_ignore_ascii_case("final_answer") {
|
||||||
|
Ok(LlmResponse::FinalAnswer {
|
||||||
|
thought,
|
||||||
|
answer: input,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
// Parse arguments as JSON, falling back to a string if invalid.
|
||||||
|
let args = serde_json::from_str(&input).unwrap_or_else(|_| json!(input));
|
||||||
|
Ok(LlmResponse::ToolCall {
|
||||||
|
thought,
|
||||||
|
tool_name: action,
|
||||||
|
arguments: args,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a single tool call via the MCP client.
|
||||||
|
async fn execute_tool(
|
||||||
|
&self,
|
||||||
|
name: &str,
|
||||||
|
arguments: serde_json::Value,
|
||||||
|
) -> CoreResult<McpToolResponse> {
|
||||||
|
// For potentially unsafe tools (write/delete) ask for UI confirmation
|
||||||
|
// if a controller is available.
|
||||||
|
let dangerous = name.contains("write") || name.contains("delete");
|
||||||
|
if dangerous {
|
||||||
|
if let Some(controller) = &self.ui_controller {
|
||||||
|
let prompt = format!(
|
||||||
|
"Confirm execution of potentially unsafe tool '{}' with args {}?",
|
||||||
|
name, arguments
|
||||||
|
);
|
||||||
|
if !controller.confirm(&prompt).await {
|
||||||
|
return Err(Error::PermissionDenied(format!(
|
||||||
|
"Tool '{}' denied by user",
|
||||||
|
name
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let call = McpToolCall {
|
||||||
|
name: name.to_string(),
|
||||||
|
arguments,
|
||||||
|
};
|
||||||
|
self.tool_client.call_tool(call).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run the full ReAct loop and return the final answer.
|
||||||
|
pub async fn run(&self, query: String) -> std::result::Result<String, AgentError> {
|
||||||
|
let tools = self.discover_tools().await.map_err(AgentError::Mcp)?;
|
||||||
|
|
||||||
|
// Build system prompt with ReAct format instructions
|
||||||
|
let tools_desc = tools
|
||||||
|
.iter()
|
||||||
|
.map(|t| {
|
||||||
|
let schema_str = serde_json::to_string_pretty(&t.input_schema)
|
||||||
|
.unwrap_or_else(|_| "{}".to_string());
|
||||||
|
format!(
|
||||||
|
"- {}: {}\n Input schema: {}",
|
||||||
|
t.name, t.description, schema_str
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n");
|
||||||
|
|
||||||
|
let system_prompt = format!(
|
||||||
|
"You are an AI assistant that uses the ReAct (Reasoning + Acting) pattern to solve tasks.\n\n\
|
||||||
|
You must ALWAYS respond in this exact format:\n\n\
|
||||||
|
THOUGHT: <your reasoning about what to do next>\n\
|
||||||
|
ACTION: <tool_name or \"final_answer\">\n\
|
||||||
|
ACTION_INPUT: <JSON arguments for the tool, or the final answer text>\n\n\
|
||||||
|
Available tools:\n{}\n\n\
|
||||||
|
HOW IT WORKS:\n\
|
||||||
|
1. When you call a tool, you will receive its output in the next message\n\
|
||||||
|
2. After receiving the tool output, analyze it and either:\n\
|
||||||
|
a) Use the information to provide a final answer\n\
|
||||||
|
b) Call another tool if you need more information\n\
|
||||||
|
3. When you have the information needed to answer the user's question, provide a final answer\n\n\
|
||||||
|
To provide a final answer:\n\
|
||||||
|
THOUGHT: <summary of what you learned>\n\
|
||||||
|
ACTION: final_answer\n\
|
||||||
|
ACTION_INPUT: <your complete answer using the information from the tools>\n\n\
|
||||||
|
IMPORTANT: You MUST follow this format exactly. Do not deviate from it.\n\
|
||||||
|
IMPORTANT: Only use the tools listed above. Do not try to use tools that are not listed.\n\
|
||||||
|
IMPORTANT: When providing the final answer, include the actual information you learned, not just the tool arguments.",
|
||||||
|
tools_desc
|
||||||
|
);
|
||||||
|
|
||||||
|
// Initialize conversation with system prompt and user query
|
||||||
|
let mut messages = vec![Message::system(system_prompt.clone()), Message::user(query)];
|
||||||
|
|
||||||
|
// Cancellation flag set when Ctrl+C is received.
|
||||||
|
let cancelled = Arc::new(AtomicBool::new(false));
|
||||||
|
let cancel_flag = cancelled.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
// Wait for Ctrl+C signal.
|
||||||
|
let _ = signal::ctrl_c().await;
|
||||||
|
cancel_flag.store(true, Ordering::SeqCst);
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut tool_calls = 0usize;
|
||||||
|
for _ in 0..self.config.max_iterations {
|
||||||
|
if cancelled.load(Ordering::SeqCst) {
|
||||||
|
return Err(AgentError::Cancelled);
|
||||||
|
}
|
||||||
|
// Build a ChatRequest for the provider.
|
||||||
|
let chat_req = ChatRequest {
|
||||||
|
model: self.config.model.clone(),
|
||||||
|
messages: messages.clone(),
|
||||||
|
parameters: crate::types::ChatParameters {
|
||||||
|
temperature: self.config.temperature,
|
||||||
|
max_tokens: self.config.max_tokens,
|
||||||
|
stream: false,
|
||||||
|
extra: Default::default(),
|
||||||
|
},
|
||||||
|
tools: Some(tools.clone()),
|
||||||
|
};
|
||||||
|
let raw_resp = self
|
||||||
|
.llm_client
|
||||||
|
.chat(chat_req)
|
||||||
|
.await
|
||||||
|
.map_err(AgentError::Provider)?;
|
||||||
|
let parsed = self
|
||||||
|
.parse_response(&raw_resp.message.content)
|
||||||
|
.map_err(|e| {
|
||||||
|
eprintln!("\n=== PARSE ERROR ===");
|
||||||
|
eprintln!("Error: {:?}", e);
|
||||||
|
eprintln!("LLM Response:\n{}", raw_resp.message.content);
|
||||||
|
eprintln!("=== END ===\n");
|
||||||
|
e
|
||||||
|
})?;
|
||||||
|
match parsed {
|
||||||
|
LlmResponse::Reasoning { thought } => {
|
||||||
|
// Append the reasoning as an assistant message.
|
||||||
|
messages.push(Message::assistant(thought));
|
||||||
|
}
|
||||||
|
LlmResponse::ToolCall {
|
||||||
|
thought,
|
||||||
|
tool_name,
|
||||||
|
arguments,
|
||||||
|
} => {
|
||||||
|
// Record the thought.
|
||||||
|
messages.push(Message::assistant(thought));
|
||||||
|
// Enforce tool call budget.
|
||||||
|
tool_calls += 1;
|
||||||
|
if tool_calls > self.config.max_tool_calls {
|
||||||
|
return Err(AgentError::MaxIterationsReached(self.config.max_iterations));
|
||||||
|
}
|
||||||
|
// Execute tool.
|
||||||
|
let args_clone = arguments.clone();
|
||||||
|
let tool_resp = self
|
||||||
|
.execute_tool(&tool_name, args_clone.clone())
|
||||||
|
.await
|
||||||
|
.map_err(AgentError::Mcp)?;
|
||||||
|
// Convert tool output to a string for the message.
|
||||||
|
let output_str = tool_resp
|
||||||
|
.output
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.to_string())
|
||||||
|
.unwrap_or_else(|| tool_resp.output.to_string());
|
||||||
|
// Audit log the tool execution.
|
||||||
|
if let Some(config_dir) = dirs::config_dir() {
|
||||||
|
let log_path = config_dir.join("owlen/logs/tool_execution.log");
|
||||||
|
if let Some(parent) = log_path.parent() {
|
||||||
|
let _ = std::fs::create_dir_all(parent);
|
||||||
|
}
|
||||||
|
if let Ok(mut file) =
|
||||||
|
OpenOptions::new().create(true).append(true).open(&log_path)
|
||||||
|
{
|
||||||
|
let ts = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs();
|
||||||
|
let _ = writeln!(
|
||||||
|
file,
|
||||||
|
"{} | tool: {} | args: {} | output: {}",
|
||||||
|
ts, tool_name, args_clone, output_str
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
messages.push(Message::tool(tool_name, output_str));
|
||||||
|
}
|
||||||
|
LlmResponse::FinalAnswer { thought, answer } => {
|
||||||
|
// Append final thought and answer, then return.
|
||||||
|
messages.push(Message::assistant(thought));
|
||||||
|
// The final answer should be a single assistant message.
|
||||||
|
messages.push(Message::assistant(answer.clone()));
|
||||||
|
return Ok(answer);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(AgentError::MaxIterationsReached(self.config.max_iterations))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -14,6 +14,9 @@ pub const DEFAULT_CONFIG_PATH: &str = "~/.config/owlen/config.toml";
|
|||||||
pub struct Config {
|
pub struct Config {
|
||||||
/// General application settings
|
/// General application settings
|
||||||
pub general: GeneralSettings,
|
pub general: GeneralSettings,
|
||||||
|
/// MCP (Multi-Client-Provider) settings
|
||||||
|
#[serde(default)]
|
||||||
|
pub mcp: McpSettings,
|
||||||
/// Provider specific configuration keyed by provider name
|
/// Provider specific configuration keyed by provider name
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub providers: HashMap<String, ProviderConfig>,
|
pub providers: HashMap<String, ProviderConfig>,
|
||||||
@@ -26,27 +29,36 @@ pub struct Config {
|
|||||||
/// Input handling preferences
|
/// Input handling preferences
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub input: InputSettings,
|
pub input: InputSettings,
|
||||||
|
/// Privacy controls for tooling and network usage
|
||||||
|
#[serde(default)]
|
||||||
|
pub privacy: PrivacySettings,
|
||||||
|
/// Security controls for sandboxing and resource limits
|
||||||
|
#[serde(default)]
|
||||||
|
pub security: SecuritySettings,
|
||||||
|
/// Per-tool configuration toggles
|
||||||
|
#[serde(default)]
|
||||||
|
pub tools: ToolSettings,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for Config {
|
impl Default for Config {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
let mut providers = HashMap::new();
|
let mut providers = HashMap::new();
|
||||||
|
providers.insert("ollama".to_string(), default_ollama_provider_config());
|
||||||
providers.insert(
|
providers.insert(
|
||||||
"ollama".to_string(),
|
"ollama-cloud".to_string(),
|
||||||
ProviderConfig {
|
default_ollama_cloud_provider_config(),
|
||||||
provider_type: "ollama".to_string(),
|
|
||||||
base_url: Some("http://localhost:11434".to_string()),
|
|
||||||
api_key: None,
|
|
||||||
extra: HashMap::new(),
|
|
||||||
},
|
|
||||||
);
|
);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
general: GeneralSettings::default(),
|
general: GeneralSettings::default(),
|
||||||
|
mcp: McpSettings::default(),
|
||||||
providers,
|
providers,
|
||||||
ui: UiSettings::default(),
|
ui: UiSettings::default(),
|
||||||
storage: StorageSettings::default(),
|
storage: StorageSettings::default(),
|
||||||
input: InputSettings::default(),
|
input: InputSettings::default(),
|
||||||
|
privacy: PrivacySettings::default(),
|
||||||
|
security: SecuritySettings::default(),
|
||||||
|
tools: ToolSettings::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -120,17 +132,26 @@ impl Config {
|
|||||||
self.general.default_provider = "ollama".to_string();
|
self.general.default_provider = "ollama".to_string();
|
||||||
}
|
}
|
||||||
|
|
||||||
if !self.providers.contains_key("ollama") {
|
ensure_provider_config(self, "ollama");
|
||||||
self.providers.insert(
|
ensure_provider_config(self, "ollama-cloud");
|
||||||
"ollama".to_string(),
|
}
|
||||||
ProviderConfig {
|
}
|
||||||
provider_type: "ollama".to_string(),
|
|
||||||
base_url: Some("http://localhost:11434".to_string()),
|
fn default_ollama_provider_config() -> ProviderConfig {
|
||||||
api_key: None,
|
ProviderConfig {
|
||||||
extra: HashMap::new(),
|
provider_type: "ollama".to_string(),
|
||||||
},
|
base_url: Some("http://localhost:11434".to_string()),
|
||||||
);
|
api_key: None,
|
||||||
}
|
extra: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_ollama_cloud_provider_config() -> ProviderConfig {
|
||||||
|
ProviderConfig {
|
||||||
|
provider_type: "ollama-cloud".to_string(),
|
||||||
|
base_url: Some("https://ollama.com".to_string()),
|
||||||
|
api_key: None,
|
||||||
|
extra: HashMap::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -185,6 +206,174 @@ impl Default for GeneralSettings {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// MCP (Multi-Client-Provider) settings
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct McpSettings {
|
||||||
|
#[serde(default)]
|
||||||
|
pub mode: McpMode,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum McpMode {
|
||||||
|
#[default]
|
||||||
|
Legacy,
|
||||||
|
Enabled,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Privacy controls governing network access and storage
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct PrivacySettings {
|
||||||
|
#[serde(default = "PrivacySettings::default_remote_search")]
|
||||||
|
pub enable_remote_search: bool,
|
||||||
|
#[serde(default)]
|
||||||
|
pub cache_web_results: bool,
|
||||||
|
#[serde(default)]
|
||||||
|
pub retain_history_days: u32,
|
||||||
|
#[serde(default = "PrivacySettings::default_require_consent")]
|
||||||
|
pub require_consent_per_session: bool,
|
||||||
|
#[serde(default = "PrivacySettings::default_encrypt_local_data")]
|
||||||
|
pub encrypt_local_data: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PrivacySettings {
|
||||||
|
const fn default_remote_search() -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn default_require_consent() -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn default_encrypt_local_data() -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for PrivacySettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enable_remote_search: Self::default_remote_search(),
|
||||||
|
cache_web_results: false,
|
||||||
|
retain_history_days: 0,
|
||||||
|
require_consent_per_session: Self::default_require_consent(),
|
||||||
|
encrypt_local_data: Self::default_encrypt_local_data(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Security settings that constrain tool execution
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct SecuritySettings {
|
||||||
|
#[serde(default = "SecuritySettings::default_enable_sandboxing")]
|
||||||
|
pub enable_sandboxing: bool,
|
||||||
|
#[serde(default = "SecuritySettings::default_timeout")]
|
||||||
|
pub sandbox_timeout_seconds: u64,
|
||||||
|
#[serde(default = "SecuritySettings::default_max_memory")]
|
||||||
|
pub max_memory_mb: u64,
|
||||||
|
#[serde(default = "SecuritySettings::default_allowed_tools")]
|
||||||
|
pub allowed_tools: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SecuritySettings {
|
||||||
|
const fn default_enable_sandboxing() -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn default_timeout() -> u64 {
|
||||||
|
30
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn default_max_memory() -> u64 {
|
||||||
|
512
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_allowed_tools() -> Vec<String> {
|
||||||
|
vec![
|
||||||
|
"web_search".to_string(),
|
||||||
|
"code_exec".to_string(),
|
||||||
|
"file_write".to_string(),
|
||||||
|
"file_delete".to_string(),
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for SecuritySettings {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enable_sandboxing: Self::default_enable_sandboxing(),
|
||||||
|
sandbox_timeout_seconds: Self::default_timeout(),
|
||||||
|
max_memory_mb: Self::default_max_memory(),
|
||||||
|
allowed_tools: Self::default_allowed_tools(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-tool configuration toggles
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct ToolSettings {
|
||||||
|
#[serde(default)]
|
||||||
|
pub web_search: WebSearchToolConfig,
|
||||||
|
#[serde(default)]
|
||||||
|
pub code_exec: CodeExecToolConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct WebSearchToolConfig {
|
||||||
|
#[serde(default)]
|
||||||
|
pub enabled: bool,
|
||||||
|
#[serde(default)]
|
||||||
|
pub api_key: String,
|
||||||
|
#[serde(default = "WebSearchToolConfig::default_max_results")]
|
||||||
|
pub max_results: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WebSearchToolConfig {
|
||||||
|
const fn default_max_results() -> u32 {
|
||||||
|
5
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for WebSearchToolConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false,
|
||||||
|
api_key: String::new(),
|
||||||
|
max_results: Self::default_max_results(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CodeExecToolConfig {
|
||||||
|
#[serde(default)]
|
||||||
|
pub enabled: bool,
|
||||||
|
#[serde(default = "CodeExecToolConfig::default_allowed_languages")]
|
||||||
|
pub allowed_languages: Vec<String>,
|
||||||
|
#[serde(default = "CodeExecToolConfig::default_timeout")]
|
||||||
|
pub timeout_seconds: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CodeExecToolConfig {
|
||||||
|
fn default_allowed_languages() -> Vec<String> {
|
||||||
|
vec!["python".to_string(), "javascript".to_string()]
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn default_timeout() -> u64 {
|
||||||
|
30
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for CodeExecToolConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
enabled: false,
|
||||||
|
allowed_languages: Self::default_allowed_languages(),
|
||||||
|
timeout_seconds: Self::default_timeout(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// UI preferences that consumers can respect as needed
|
/// UI preferences that consumers can respect as needed
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct UiSettings {
|
pub struct UiSettings {
|
||||||
@@ -343,15 +532,32 @@ impl Default for InputSettings {
|
|||||||
|
|
||||||
/// Convenience accessor for an Ollama provider entry, creating a default if missing
|
/// Convenience accessor for an Ollama provider entry, creating a default if missing
|
||||||
pub fn ensure_ollama_config(config: &mut Config) -> &ProviderConfig {
|
pub fn ensure_ollama_config(config: &mut Config) -> &ProviderConfig {
|
||||||
config
|
ensure_provider_config(config, "ollama")
|
||||||
.providers
|
}
|
||||||
.entry("ollama".to_string())
|
|
||||||
.or_insert_with(|| ProviderConfig {
|
/// Ensure a provider configuration exists for the requested provider name
|
||||||
provider_type: "ollama".to_string(),
|
pub fn ensure_provider_config<'a>(
|
||||||
base_url: Some("http://localhost:11434".to_string()),
|
config: &'a mut Config,
|
||||||
api_key: None,
|
provider_name: &str,
|
||||||
extra: HashMap::new(),
|
) -> &'a ProviderConfig {
|
||||||
})
|
use std::collections::hash_map::Entry;
|
||||||
|
|
||||||
|
match config.providers.entry(provider_name.to_string()) {
|
||||||
|
Entry::Occupied(entry) => entry.into_mut(),
|
||||||
|
Entry::Vacant(entry) => {
|
||||||
|
let default = match provider_name {
|
||||||
|
"ollama-cloud" => default_ollama_cloud_provider_config(),
|
||||||
|
"ollama" => default_ollama_provider_config(),
|
||||||
|
other => ProviderConfig {
|
||||||
|
provider_type: other.to_string(),
|
||||||
|
base_url: None,
|
||||||
|
api_key: None,
|
||||||
|
extra: HashMap::new(),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
entry.insert(default)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Calculate absolute timeout for session data based on configuration
|
/// Calculate absolute timeout for session data based on configuration
|
||||||
@@ -404,4 +610,21 @@ mod tests {
|
|||||||
let path = config.storage.conversation_path();
|
let path = config.storage.conversation_path();
|
||||||
assert!(path.to_string_lossy().contains("custom/path"));
|
assert!(path.to_string_lossy().contains("custom/path"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_config_contains_local_and_cloud_providers() {
|
||||||
|
let config = Config::default();
|
||||||
|
assert!(config.providers.contains_key("ollama"));
|
||||||
|
assert!(config.providers.contains_key("ollama-cloud"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ensure_provider_config_backfills_cloud_defaults() {
|
||||||
|
let mut config = Config::default();
|
||||||
|
config.providers.remove("ollama-cloud");
|
||||||
|
|
||||||
|
let cloud = ensure_provider_config(&mut config, "ollama-cloud");
|
||||||
|
assert_eq!(cloud.provider_type, "ollama-cloud");
|
||||||
|
assert_eq!(cloud.base_url.as_deref(), Some("https://ollama.com"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
295
crates/owlen-core/src/consent.rs
Normal file
295
crates/owlen-core/src/consent.rs
Normal file
@@ -0,0 +1,295 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::io::{self, Write};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
use crate::encryption::VaultHandle;
|
||||||
|
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct ConsentRequest {
|
||||||
|
pub tool_name: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Scope of consent grant
|
||||||
|
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
|
||||||
|
pub enum ConsentScope {
|
||||||
|
/// Grant only for this single operation
|
||||||
|
Once,
|
||||||
|
/// Grant for the duration of the current session
|
||||||
|
Session,
|
||||||
|
/// Grant permanently (persisted across sessions)
|
||||||
|
Permanent,
|
||||||
|
/// Explicitly denied
|
||||||
|
Denied,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||||
|
pub struct ConsentRecord {
|
||||||
|
pub tool_name: String,
|
||||||
|
pub scope: ConsentScope,
|
||||||
|
pub timestamp: DateTime<Utc>,
|
||||||
|
pub data_types: Vec<String>,
|
||||||
|
pub external_endpoints: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize, Default)]
|
||||||
|
pub struct ConsentManager {
|
||||||
|
/// Permanent consent records (persisted to vault)
|
||||||
|
permanent_records: HashMap<String, ConsentRecord>,
|
||||||
|
/// Session-scoped consent (cleared on manager drop or explicit clear)
|
||||||
|
#[serde(skip)]
|
||||||
|
session_records: HashMap<String, ConsentRecord>,
|
||||||
|
/// Once-scoped consent (used once then cleared)
|
||||||
|
#[serde(skip)]
|
||||||
|
once_records: HashMap<String, ConsentRecord>,
|
||||||
|
/// Pending consent requests (to prevent duplicate prompts)
|
||||||
|
#[serde(skip)]
|
||||||
|
pending_requests: HashMap<String, ()>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ConsentManager {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load consent records from vault storage
|
||||||
|
pub fn from_vault(vault: &Arc<std::sync::Mutex<VaultHandle>>) -> Self {
|
||||||
|
let guard = vault.lock().expect("Vault mutex poisoned");
|
||||||
|
if let Some(consent_data) = guard.settings().get("consent_records") {
|
||||||
|
if let Ok(permanent_records) =
|
||||||
|
serde_json::from_value::<HashMap<String, ConsentRecord>>(consent_data.clone())
|
||||||
|
{
|
||||||
|
return Self {
|
||||||
|
permanent_records,
|
||||||
|
session_records: HashMap::new(),
|
||||||
|
once_records: HashMap::new(),
|
||||||
|
pending_requests: HashMap::new(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Persist permanent consent records to vault storage
|
||||||
|
pub fn persist_to_vault(&self, vault: &Arc<std::sync::Mutex<VaultHandle>>) -> Result<()> {
|
||||||
|
let mut guard = vault.lock().expect("Vault mutex poisoned");
|
||||||
|
let consent_json = serde_json::to_value(&self.permanent_records)?;
|
||||||
|
guard
|
||||||
|
.settings_mut()
|
||||||
|
.insert("consent_records".to_string(), consent_json);
|
||||||
|
guard.persist()?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn request_consent(
|
||||||
|
&mut self,
|
||||||
|
tool_name: &str,
|
||||||
|
data_types: Vec<String>,
|
||||||
|
endpoints: Vec<String>,
|
||||||
|
) -> Result<ConsentScope> {
|
||||||
|
// Check if already granted permanently
|
||||||
|
if let Some(existing) = self.permanent_records.get(tool_name) {
|
||||||
|
if existing.scope == ConsentScope::Permanent {
|
||||||
|
return Ok(ConsentScope::Permanent);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if granted for session
|
||||||
|
if let Some(existing) = self.session_records.get(tool_name) {
|
||||||
|
if existing.scope == ConsentScope::Session {
|
||||||
|
return Ok(ConsentScope::Session);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if request is already pending (prevent duplicate prompts)
|
||||||
|
if self.pending_requests.contains_key(tool_name) {
|
||||||
|
// Wait for the other prompt to complete by returning denied temporarily
|
||||||
|
// The caller should retry after a short delay
|
||||||
|
return Ok(ConsentScope::Denied);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Mark as pending
|
||||||
|
self.pending_requests.insert(tool_name.to_string(), ());
|
||||||
|
|
||||||
|
// Show consent dialog and get scope
|
||||||
|
let scope = self.show_consent_dialog(tool_name, &data_types, &endpoints)?;
|
||||||
|
|
||||||
|
// Remove from pending
|
||||||
|
self.pending_requests.remove(tool_name);
|
||||||
|
|
||||||
|
// Create record based on scope
|
||||||
|
let record = ConsentRecord {
|
||||||
|
tool_name: tool_name.to_string(),
|
||||||
|
scope: scope.clone(),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
data_types,
|
||||||
|
external_endpoints: endpoints,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Store in appropriate location
|
||||||
|
match scope {
|
||||||
|
ConsentScope::Permanent => {
|
||||||
|
self.permanent_records.insert(tool_name.to_string(), record);
|
||||||
|
}
|
||||||
|
ConsentScope::Session => {
|
||||||
|
self.session_records.insert(tool_name.to_string(), record);
|
||||||
|
}
|
||||||
|
ConsentScope::Once | ConsentScope::Denied => {
|
||||||
|
// Don't store, just return the decision
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(scope)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Grant consent programmatically (for TUI or automated flows)
|
||||||
|
pub fn grant_consent(
|
||||||
|
&mut self,
|
||||||
|
tool_name: &str,
|
||||||
|
data_types: Vec<String>,
|
||||||
|
endpoints: Vec<String>,
|
||||||
|
) {
|
||||||
|
self.grant_consent_with_scope(tool_name, data_types, endpoints, ConsentScope::Permanent);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Grant consent with specific scope
|
||||||
|
pub fn grant_consent_with_scope(
|
||||||
|
&mut self,
|
||||||
|
tool_name: &str,
|
||||||
|
data_types: Vec<String>,
|
||||||
|
endpoints: Vec<String>,
|
||||||
|
scope: ConsentScope,
|
||||||
|
) {
|
||||||
|
let record = ConsentRecord {
|
||||||
|
tool_name: tool_name.to_string(),
|
||||||
|
scope: scope.clone(),
|
||||||
|
timestamp: Utc::now(),
|
||||||
|
data_types,
|
||||||
|
external_endpoints: endpoints,
|
||||||
|
};
|
||||||
|
|
||||||
|
match scope {
|
||||||
|
ConsentScope::Permanent => {
|
||||||
|
self.permanent_records.insert(tool_name.to_string(), record);
|
||||||
|
}
|
||||||
|
ConsentScope::Session => {
|
||||||
|
self.session_records.insert(tool_name.to_string(), record);
|
||||||
|
}
|
||||||
|
ConsentScope::Once => {
|
||||||
|
self.once_records.insert(tool_name.to_string(), record);
|
||||||
|
}
|
||||||
|
ConsentScope::Denied => {} // Denied is not stored
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if consent is needed (returns None if already granted, Some(info) if needed)
|
||||||
|
pub fn check_consent_needed(&self, tool_name: &str) -> Option<ConsentRequest> {
|
||||||
|
if self.has_consent(tool_name) {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(ConsentRequest {
|
||||||
|
tool_name: tool_name.to_string(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn has_consent(&self, tool_name: &str) -> bool {
|
||||||
|
// Check permanent first, then session, then once
|
||||||
|
self.permanent_records
|
||||||
|
.get(tool_name)
|
||||||
|
.map(|r| r.scope == ConsentScope::Permanent)
|
||||||
|
.or_else(|| {
|
||||||
|
self.session_records
|
||||||
|
.get(tool_name)
|
||||||
|
.map(|r| r.scope == ConsentScope::Session)
|
||||||
|
})
|
||||||
|
.or_else(|| {
|
||||||
|
self.once_records
|
||||||
|
.get(tool_name)
|
||||||
|
.map(|r| r.scope == ConsentScope::Once)
|
||||||
|
})
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Consume "once" consent for a tool (clears it after first use)
|
||||||
|
pub fn consume_once_consent(&mut self, tool_name: &str) {
|
||||||
|
self.once_records.remove(tool_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn revoke_consent(&mut self, tool_name: &str) {
|
||||||
|
self.permanent_records.remove(tool_name);
|
||||||
|
self.session_records.remove(tool_name);
|
||||||
|
self.once_records.remove(tool_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn clear_all_consent(&mut self) {
|
||||||
|
self.permanent_records.clear();
|
||||||
|
self.session_records.clear();
|
||||||
|
self.once_records.clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clear only session-scoped consent (useful when starting new session)
|
||||||
|
pub fn clear_session_consent(&mut self) {
|
||||||
|
self.session_records.clear();
|
||||||
|
self.once_records.clear(); // Also clear once consent on session clear
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if consent is needed for a tool (non-blocking)
|
||||||
|
/// Returns Some with consent details if needed, None if already granted
|
||||||
|
pub fn check_if_consent_needed(
|
||||||
|
&self,
|
||||||
|
tool_name: &str,
|
||||||
|
data_types: Vec<String>,
|
||||||
|
endpoints: Vec<String>,
|
||||||
|
) -> Option<(String, Vec<String>, Vec<String>)> {
|
||||||
|
if self.has_consent(tool_name) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some((tool_name.to_string(), data_types, endpoints))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn show_consent_dialog(
|
||||||
|
&self,
|
||||||
|
tool_name: &str,
|
||||||
|
data_types: &[String],
|
||||||
|
endpoints: &[String],
|
||||||
|
) -> Result<ConsentScope> {
|
||||||
|
// TEMPORARY: Auto-grant session consent when not in a proper terminal (TUI mode)
|
||||||
|
// TODO: Integrate consent UI into the TUI event loop
|
||||||
|
use std::io::IsTerminal;
|
||||||
|
if !io::stdin().is_terminal() || std::env::var("OWLEN_AUTO_CONSENT").is_ok() {
|
||||||
|
eprintln!("Auto-granting session consent for {} (TUI mode)", tool_name);
|
||||||
|
return Ok(ConsentScope::Session);
|
||||||
|
}
|
||||||
|
|
||||||
|
println!("\n╔══════════════════════════════════════════════════╗");
|
||||||
|
println!("║ 🔒 PRIVACY CONSENT REQUIRED 🔒 ║");
|
||||||
|
println!("╚══════════════════════════════════════════════════╝");
|
||||||
|
println!();
|
||||||
|
println!("Tool: {}", tool_name);
|
||||||
|
println!("Data: {}", data_types.join(", "));
|
||||||
|
println!("Endpoints: {}", endpoints.join(", "));
|
||||||
|
println!();
|
||||||
|
println!("Choose consent scope:");
|
||||||
|
println!(" [1] Allow once - Grant only for this operation");
|
||||||
|
println!(" [2] Allow session - Grant for current session");
|
||||||
|
println!(" [3] Allow always - Grant permanently");
|
||||||
|
println!(" [4] Deny - Reject this operation");
|
||||||
|
println!();
|
||||||
|
print!("Enter choice (1-4) [default: 4]: ");
|
||||||
|
io::stdout().flush()?;
|
||||||
|
|
||||||
|
let mut input = String::new();
|
||||||
|
io::stdin().read_line(&mut input)?;
|
||||||
|
|
||||||
|
match input.trim() {
|
||||||
|
"1" => Ok(ConsentScope::Once),
|
||||||
|
"2" => Ok(ConsentScope::Session),
|
||||||
|
"3" => Ok(ConsentScope::Permanent),
|
||||||
|
_ => Ok(ConsentScope::Denied),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -3,7 +3,6 @@ use crate::types::{Conversation, Message};
|
|||||||
use crate::Result;
|
use crate::Result;
|
||||||
use serde_json::{Number, Value};
|
use serde_json::{Number, Value};
|
||||||
use std::collections::{HashMap, VecDeque};
|
use std::collections::{HashMap, VecDeque};
|
||||||
use std::path::{Path, PathBuf};
|
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
@@ -214,6 +213,25 @@ impl ConversationManager {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Set tool calls on a streaming message
|
||||||
|
pub fn set_tool_calls_on_message(
|
||||||
|
&mut self,
|
||||||
|
message_id: Uuid,
|
||||||
|
tool_calls: Vec<crate::types::ToolCall>,
|
||||||
|
) -> Result<()> {
|
||||||
|
let index = self
|
||||||
|
.message_index
|
||||||
|
.get(&message_id)
|
||||||
|
.copied()
|
||||||
|
.ok_or_else(|| crate::Error::Unknown(format!("Unknown message id: {message_id}")))?;
|
||||||
|
|
||||||
|
if let Some(message) = self.active_mut().messages.get_mut(index) {
|
||||||
|
message.tool_calls = Some(tool_calls);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
/// Update the active model (used when user changes model mid session)
|
/// Update the active model (used when user changes model mid session)
|
||||||
pub fn set_model(&mut self, model: impl Into<String>) {
|
pub fn set_model(&mut self, model: impl Into<String>) {
|
||||||
self.active.model = model.into();
|
self.active.model = model.into();
|
||||||
@@ -268,36 +286,40 @@ impl ConversationManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Save the active conversation to disk
|
/// Save the active conversation to disk
|
||||||
pub fn save_active(&self, storage: &StorageManager, name: Option<String>) -> Result<PathBuf> {
|
pub async fn save_active(
|
||||||
storage.save_conversation(&self.active, name)
|
&self,
|
||||||
|
storage: &StorageManager,
|
||||||
|
name: Option<String>,
|
||||||
|
) -> Result<Uuid> {
|
||||||
|
storage.save_conversation(&self.active, name).await?;
|
||||||
|
Ok(self.active.id)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Save the active conversation to disk with a description
|
/// Save the active conversation to disk with a description
|
||||||
pub fn save_active_with_description(
|
pub async fn save_active_with_description(
|
||||||
&self,
|
&self,
|
||||||
storage: &StorageManager,
|
storage: &StorageManager,
|
||||||
name: Option<String>,
|
name: Option<String>,
|
||||||
description: Option<String>,
|
description: Option<String>,
|
||||||
) -> Result<PathBuf> {
|
) -> Result<Uuid> {
|
||||||
storage.save_conversation_with_description(&self.active, name, description)
|
storage
|
||||||
|
.save_conversation_with_description(&self.active, name, description)
|
||||||
|
.await?;
|
||||||
|
Ok(self.active.id)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Load a conversation from disk and make it active
|
/// Load a conversation from storage and make it active
|
||||||
pub fn load_from_disk(
|
pub async fn load_saved(&mut self, storage: &StorageManager, id: Uuid) -> Result<()> {
|
||||||
&mut self,
|
let conversation = storage.load_conversation(id).await?;
|
||||||
storage: &StorageManager,
|
|
||||||
path: impl AsRef<Path>,
|
|
||||||
) -> Result<()> {
|
|
||||||
let conversation = storage.load_conversation(path)?;
|
|
||||||
self.load(conversation);
|
self.load(conversation);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// List all saved sessions
|
/// List all saved sessions
|
||||||
pub fn list_saved_sessions(
|
pub async fn list_saved_sessions(
|
||||||
storage: &StorageManager,
|
storage: &StorageManager,
|
||||||
) -> Result<Vec<crate::storage::SessionMeta>> {
|
) -> Result<Vec<crate::storage::SessionMeta>> {
|
||||||
storage.list_sessions()
|
storage.list_sessions().await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
69
crates/owlen-core/src/credentials.rs
Normal file
69
crates/owlen-core/src/credentials.rs
Normal file
@@ -0,0 +1,69 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
use crate::{storage::StorageManager, Error, Result};
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize, Debug)]
|
||||||
|
pub struct ApiCredentials {
|
||||||
|
pub api_key: String,
|
||||||
|
pub endpoint: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct CredentialManager {
|
||||||
|
storage: Arc<StorageManager>,
|
||||||
|
master_key: Arc<Vec<u8>>,
|
||||||
|
namespace: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CredentialManager {
|
||||||
|
pub fn new(storage: Arc<StorageManager>, master_key: Arc<Vec<u8>>) -> Self {
|
||||||
|
Self {
|
||||||
|
storage,
|
||||||
|
master_key,
|
||||||
|
namespace: "owlen".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn namespaced_key(&self, tool_name: &str) -> String {
|
||||||
|
format!("{}_{}", self.namespace, tool_name)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn store_credentials(
|
||||||
|
&self,
|
||||||
|
tool_name: &str,
|
||||||
|
credentials: &ApiCredentials,
|
||||||
|
) -> Result<()> {
|
||||||
|
let key = self.namespaced_key(tool_name);
|
||||||
|
let payload = serde_json::to_vec(credentials).map_err(|e| {
|
||||||
|
Error::Storage(format!(
|
||||||
|
"Failed to serialize credentials for secure storage: {e}"
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
self.storage
|
||||||
|
.store_secure_item(&key, &payload, &self.master_key)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_credentials(&self, tool_name: &str) -> Result<Option<ApiCredentials>> {
|
||||||
|
let key = self.namespaced_key(tool_name);
|
||||||
|
match self
|
||||||
|
.storage
|
||||||
|
.load_secure_item(&key, &self.master_key)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
Some(bytes) => {
|
||||||
|
let creds = serde_json::from_slice(&bytes).map_err(|e| {
|
||||||
|
Error::Storage(format!("Failed to deserialize stored credentials: {e}"))
|
||||||
|
})?;
|
||||||
|
Ok(Some(creds))
|
||||||
|
}
|
||||||
|
None => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete_credentials(&self, tool_name: &str) -> Result<()> {
|
||||||
|
let key = self.namespaced_key(tool_name);
|
||||||
|
self.storage.delete_secure_item(&key).await
|
||||||
|
}
|
||||||
|
}
|
||||||
241
crates/owlen-core/src/encryption.rs
Normal file
241
crates/owlen-core/src/encryption.rs
Normal file
@@ -0,0 +1,241 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::fs;
|
||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
use aes_gcm::{
|
||||||
|
aead::{Aead, KeyInit},
|
||||||
|
Aes256Gcm, Nonce,
|
||||||
|
};
|
||||||
|
use anyhow::{bail, Context, Result};
|
||||||
|
use ring::digest;
|
||||||
|
use ring::rand::{SecureRandom, SystemRandom};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value as JsonValue;
|
||||||
|
|
||||||
|
pub struct EncryptedStorage {
|
||||||
|
cipher: Aes256Gcm,
|
||||||
|
storage_path: PathBuf,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize)]
|
||||||
|
struct EncryptedData {
|
||||||
|
nonce: [u8; 12],
|
||||||
|
ciphertext: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct VaultData {
|
||||||
|
pub master_key: Vec<u8>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub settings: HashMap<String, JsonValue>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct VaultHandle {
|
||||||
|
storage: EncryptedStorage,
|
||||||
|
pub data: VaultData,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VaultHandle {
|
||||||
|
pub fn master_key(&self) -> &[u8] {
|
||||||
|
&self.data.master_key
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn settings(&self) -> &HashMap<String, JsonValue> {
|
||||||
|
&self.data.settings
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn settings_mut(&mut self) -> &mut HashMap<String, JsonValue> {
|
||||||
|
&mut self.data.settings
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn persist(&self) -> Result<()> {
|
||||||
|
self.storage.store(&self.data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EncryptedStorage {
|
||||||
|
pub fn new(storage_path: PathBuf, password: &str) -> Result<Self> {
|
||||||
|
let digest = digest::digest(&digest::SHA256, password.as_bytes());
|
||||||
|
let cipher = Aes256Gcm::new_from_slice(digest.as_ref())
|
||||||
|
.map_err(|_| anyhow::anyhow!("Invalid key length for AES-256"))?;
|
||||||
|
|
||||||
|
if let Some(parent) = storage_path.parent() {
|
||||||
|
fs::create_dir_all(parent).context("Failed to ensure storage directory exists")?;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
cipher,
|
||||||
|
storage_path,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn store<T: Serialize>(&self, data: &T) -> Result<()> {
|
||||||
|
let json = serde_json::to_vec(data).context("Failed to serialize data")?;
|
||||||
|
|
||||||
|
let nonce = generate_nonce()?;
|
||||||
|
let nonce_ref = Nonce::from_slice(&nonce);
|
||||||
|
|
||||||
|
let ciphertext = self
|
||||||
|
.cipher
|
||||||
|
.encrypt(nonce_ref, json.as_ref())
|
||||||
|
.map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;
|
||||||
|
|
||||||
|
let encrypted_data = EncryptedData { nonce, ciphertext };
|
||||||
|
let encrypted_json = serde_json::to_vec(&encrypted_data)?;
|
||||||
|
|
||||||
|
fs::write(&self.storage_path, encrypted_json).context("Failed to write encrypted data")?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn load<T: for<'de> Deserialize<'de>>(&self) -> Result<T> {
|
||||||
|
let encrypted_json =
|
||||||
|
fs::read(&self.storage_path).context("Failed to read encrypted data")?;
|
||||||
|
|
||||||
|
let encrypted_data: EncryptedData =
|
||||||
|
serde_json::from_slice(&encrypted_json).context("Failed to parse encrypted data")?;
|
||||||
|
|
||||||
|
let nonce_ref = Nonce::from_slice(&encrypted_data.nonce);
|
||||||
|
let plaintext = self
|
||||||
|
.cipher
|
||||||
|
.decrypt(nonce_ref, encrypted_data.ciphertext.as_ref())
|
||||||
|
.map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))?;
|
||||||
|
|
||||||
|
let data: T =
|
||||||
|
serde_json::from_slice(&plaintext).context("Failed to deserialize decrypted data")?;
|
||||||
|
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn exists(&self) -> bool {
|
||||||
|
self.storage_path.exists()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn delete(&self) -> Result<()> {
|
||||||
|
if self.exists() {
|
||||||
|
fs::remove_file(&self.storage_path).context("Failed to delete encrypted storage")?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn verify_password(&self) -> Result<()> {
|
||||||
|
if !self.exists() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let encrypted_json =
|
||||||
|
fs::read(&self.storage_path).context("Failed to read encrypted data")?;
|
||||||
|
|
||||||
|
if encrypted_json.is_empty() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let encrypted_data: EncryptedData =
|
||||||
|
serde_json::from_slice(&encrypted_json).context("Failed to parse encrypted data")?;
|
||||||
|
|
||||||
|
let nonce_ref = Nonce::from_slice(&encrypted_data.nonce);
|
||||||
|
self.cipher
|
||||||
|
.decrypt(nonce_ref, encrypted_data.ciphertext.as_ref())
|
||||||
|
.map(|_| ())
|
||||||
|
.map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn prompt_password(prompt: &str) -> Result<String> {
|
||||||
|
let password = rpassword::prompt_password(prompt)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Failed to read password: {e}"))?;
|
||||||
|
if password.is_empty() {
|
||||||
|
bail!("Password cannot be empty");
|
||||||
|
}
|
||||||
|
Ok(password)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn prompt_new_password() -> Result<String> {
|
||||||
|
loop {
|
||||||
|
let first = prompt_password("Enter new master password: ")?;
|
||||||
|
let confirm = prompt_password("Confirm master password: ")?;
|
||||||
|
if first == confirm {
|
||||||
|
return Ok(first);
|
||||||
|
}
|
||||||
|
println!("Passwords did not match. Please try again.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn unlock_with_password(storage_path: PathBuf, password: &str) -> Result<VaultHandle> {
|
||||||
|
let storage = EncryptedStorage::new(storage_path, password)?;
|
||||||
|
let data = load_or_initialize_vault(&storage)?;
|
||||||
|
Ok(VaultHandle { storage, data })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn unlock_interactive(storage_path: PathBuf) -> Result<VaultHandle> {
|
||||||
|
if storage_path.exists() {
|
||||||
|
for attempt in 0..3 {
|
||||||
|
let password = prompt_password("Enter master password: ")?;
|
||||||
|
match unlock_with_password(storage_path.clone(), &password) {
|
||||||
|
Ok(handle) => return Ok(handle),
|
||||||
|
Err(err) => {
|
||||||
|
println!("Failed to unlock vault: {err}");
|
||||||
|
if attempt == 2 {
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
bail!("Failed to unlock encrypted storage after multiple attempts");
|
||||||
|
} else {
|
||||||
|
println!(
|
||||||
|
"No encrypted storage found at {}. Initializing a new vault.",
|
||||||
|
storage_path.display()
|
||||||
|
);
|
||||||
|
let password = prompt_new_password()?;
|
||||||
|
let storage = EncryptedStorage::new(storage_path, &password)?;
|
||||||
|
let data = VaultData {
|
||||||
|
master_key: generate_master_key()?,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
storage.store(&data)?;
|
||||||
|
Ok(VaultHandle { storage, data })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_or_initialize_vault(storage: &EncryptedStorage) -> Result<VaultData> {
|
||||||
|
match storage.load::<VaultData>() {
|
||||||
|
Ok(data) => {
|
||||||
|
if data.master_key.len() != 32 {
|
||||||
|
bail!(
|
||||||
|
"Corrupted vault: master key has invalid length ({}). \
|
||||||
|
Expected 32 bytes for AES-256. Vault cannot be recovered.",
|
||||||
|
data.master_key.len()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
if storage.exists() {
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
let data = VaultData {
|
||||||
|
master_key: generate_master_key()?,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
storage.store(&data)?;
|
||||||
|
Ok(data)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn generate_master_key() -> Result<Vec<u8>> {
|
||||||
|
let mut key = vec![0u8; 32];
|
||||||
|
SystemRandom::new()
|
||||||
|
.fill(&mut key)
|
||||||
|
.map_err(|_| anyhow::anyhow!("Failed to generate master key"))?;
|
||||||
|
Ok(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn generate_nonce() -> Result<[u8; 12]> {
|
||||||
|
let mut nonce = [0u8; 12];
|
||||||
|
let rng = SystemRandom::new();
|
||||||
|
rng.fill(&mut nonce)
|
||||||
|
.map_err(|_| anyhow::anyhow!("Failed to generate nonce"))?;
|
||||||
|
Ok(nonce)
|
||||||
|
}
|
||||||
@@ -91,6 +91,11 @@ impl MessageFormatter {
|
|||||||
Some(thinking)
|
Some(thinking)
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// If the result is empty but we have thinking content, show a placeholder
|
||||||
|
if result.trim().is_empty() && thinking_result.is_some() {
|
||||||
|
result.push_str("[Thinking...]");
|
||||||
|
}
|
||||||
|
|
||||||
(result, thinking_result)
|
(result, thinking_result)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,29 +3,45 @@
|
|||||||
//! This crate provides the foundational abstractions for building
|
//! This crate provides the foundational abstractions for building
|
||||||
//! LLM providers, routers, and MCP (Model Context Protocol) adapters.
|
//! LLM providers, routers, and MCP (Model Context Protocol) adapters.
|
||||||
|
|
||||||
|
pub mod agent;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
|
pub mod consent;
|
||||||
pub mod conversation;
|
pub mod conversation;
|
||||||
|
pub mod credentials;
|
||||||
|
pub mod encryption;
|
||||||
pub mod formatting;
|
pub mod formatting;
|
||||||
pub mod input;
|
pub mod input;
|
||||||
|
pub mod mcp;
|
||||||
pub mod model;
|
pub mod model;
|
||||||
pub mod provider;
|
pub mod provider;
|
||||||
pub mod router;
|
pub mod router;
|
||||||
|
pub mod sandbox;
|
||||||
pub mod session;
|
pub mod session;
|
||||||
pub mod storage;
|
pub mod storage;
|
||||||
pub mod theme;
|
pub mod theme;
|
||||||
|
pub mod tools;
|
||||||
pub mod types;
|
pub mod types;
|
||||||
pub mod ui;
|
pub mod ui;
|
||||||
|
pub mod validation;
|
||||||
pub mod wrap_cursor;
|
pub mod wrap_cursor;
|
||||||
|
|
||||||
|
pub use agent::*;
|
||||||
pub use config::*;
|
pub use config::*;
|
||||||
|
pub use consent::*;
|
||||||
pub use conversation::*;
|
pub use conversation::*;
|
||||||
|
pub use credentials::*;
|
||||||
|
pub use encryption::*;
|
||||||
pub use formatting::*;
|
pub use formatting::*;
|
||||||
pub use input::*;
|
pub use input::*;
|
||||||
|
pub use mcp::*;
|
||||||
pub use model::*;
|
pub use model::*;
|
||||||
pub use provider::*;
|
pub use provider::*;
|
||||||
pub use router::*;
|
pub use router::*;
|
||||||
|
pub use sandbox::*;
|
||||||
pub use session::*;
|
pub use session::*;
|
||||||
pub use theme::*;
|
pub use theme::*;
|
||||||
|
pub use tools::*;
|
||||||
|
pub use validation::*;
|
||||||
|
|
||||||
/// Result type used throughout the OWLEN ecosystem
|
/// Result type used throughout the OWLEN ecosystem
|
||||||
pub type Result<T> = std::result::Result<T, Error>;
|
pub type Result<T> = std::result::Result<T, Error>;
|
||||||
@@ -62,4 +78,10 @@ pub enum Error {
|
|||||||
|
|
||||||
#[error("Unknown error: {0}")]
|
#[error("Unknown error: {0}")]
|
||||||
Unknown(String),
|
Unknown(String),
|
||||||
|
|
||||||
|
#[error("Not implemented: {0}")]
|
||||||
|
NotImplemented(String),
|
||||||
|
|
||||||
|
#[error("Permission denied: {0}")]
|
||||||
|
PermissionDenied(String),
|
||||||
}
|
}
|
||||||
|
|||||||
113
crates/owlen-core/src/mcp.rs
Normal file
113
crates/owlen-core/src/mcp.rs
Normal file
@@ -0,0 +1,113 @@
|
|||||||
|
use crate::tools::registry::ToolRegistry;
|
||||||
|
use crate::validation::SchemaValidator;
|
||||||
|
use crate::Result;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
pub use client::McpClient;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
pub mod client;
|
||||||
|
pub mod factory;
|
||||||
|
pub mod permission;
|
||||||
|
pub mod protocol;
|
||||||
|
pub mod remote_client;
|
||||||
|
|
||||||
|
/// Descriptor for a tool exposed over MCP
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct McpToolDescriptor {
|
||||||
|
pub name: String,
|
||||||
|
pub description: String,
|
||||||
|
pub input_schema: Value,
|
||||||
|
pub requires_network: bool,
|
||||||
|
pub requires_filesystem: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Invocation payload for a tool call
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct McpToolCall {
|
||||||
|
pub name: String,
|
||||||
|
pub arguments: Value,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result returned by a tool invocation
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct McpToolResponse {
|
||||||
|
pub name: String,
|
||||||
|
pub success: bool,
|
||||||
|
pub output: Value,
|
||||||
|
pub metadata: HashMap<String, String>,
|
||||||
|
pub duration_ms: u128,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Thin MCP server facade over the tool registry
|
||||||
|
pub struct McpServer {
|
||||||
|
registry: Arc<ToolRegistry>,
|
||||||
|
validator: Arc<SchemaValidator>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl McpServer {
|
||||||
|
pub fn new(registry: Arc<ToolRegistry>, validator: Arc<SchemaValidator>) -> Self {
|
||||||
|
Self {
|
||||||
|
registry,
|
||||||
|
validator,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Enumerate the registered tools as MCP descriptors
|
||||||
|
pub fn list_tools(&self) -> Vec<McpToolDescriptor> {
|
||||||
|
self.registry
|
||||||
|
.all()
|
||||||
|
.into_iter()
|
||||||
|
.map(|tool| McpToolDescriptor {
|
||||||
|
name: tool.name().to_string(),
|
||||||
|
description: tool.description().to_string(),
|
||||||
|
input_schema: tool.schema(),
|
||||||
|
requires_network: tool.requires_network(),
|
||||||
|
requires_filesystem: tool.requires_filesystem(),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a tool call after validating inputs against the registered schema
|
||||||
|
pub async fn call_tool(&self, call: McpToolCall) -> Result<McpToolResponse> {
|
||||||
|
self.validator.validate(&call.name, &call.arguments)?;
|
||||||
|
let result = self.registry.execute(&call.name, call.arguments).await?;
|
||||||
|
Ok(McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: result.success,
|
||||||
|
output: result.output,
|
||||||
|
metadata: result.metadata,
|
||||||
|
duration_ms: duration_to_millis(result.duration),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn duration_to_millis(duration: Duration) -> u128 {
|
||||||
|
duration.as_secs() as u128 * 1_000 + u128::from(duration.subsec_millis())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct LocalMcpClient {
|
||||||
|
server: McpServer,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LocalMcpClient {
|
||||||
|
pub fn new(registry: Arc<ToolRegistry>, validator: Arc<SchemaValidator>) -> Self {
|
||||||
|
Self {
|
||||||
|
server: McpServer::new(registry, validator),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpClient for LocalMcpClient {
|
||||||
|
async fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> {
|
||||||
|
Ok(self.server.list_tools())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tool(&self, call: McpToolCall) -> Result<McpToolResponse> {
|
||||||
|
self.server.call_tool(call).await
|
||||||
|
}
|
||||||
|
}
|
||||||
51
crates/owlen-core/src/mcp/client.rs
Normal file
51
crates/owlen-core/src/mcp/client.rs
Normal file
@@ -0,0 +1,51 @@
|
|||||||
|
use super::{McpToolCall, McpToolDescriptor, McpToolResponse};
|
||||||
|
use crate::{Error, Result};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
|
||||||
|
/// Trait for a client that can interact with an MCP server
|
||||||
|
#[async_trait]
|
||||||
|
pub trait McpClient: Send + Sync {
|
||||||
|
/// List the tools available on the server
|
||||||
|
async fn list_tools(&self) -> Result<Vec<McpToolDescriptor>>;
|
||||||
|
|
||||||
|
/// Call a tool on the server
|
||||||
|
async fn call_tool(&self, call: McpToolCall) -> Result<McpToolResponse>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Placeholder for a client that connects to a remote MCP server.
|
||||||
|
pub struct RemoteMcpClient;
|
||||||
|
|
||||||
|
impl RemoteMcpClient {
|
||||||
|
pub fn new() -> Result<Self> {
|
||||||
|
// Attempt to spawn the MCP server binary located at ./target/debug/owlen-mcp-server
|
||||||
|
// The server runs over STDIO and will be managed by the client instance.
|
||||||
|
// For now we just verify that the binary exists; the actual process handling
|
||||||
|
// is performed lazily in the async methods.
|
||||||
|
let path = "./target/debug/owlen-mcp-server";
|
||||||
|
if std::path::Path::new(path).exists() {
|
||||||
|
Ok(Self)
|
||||||
|
} else {
|
||||||
|
Err(Error::NotImplemented(format!(
|
||||||
|
"Remote MCP server binary not found at {}",
|
||||||
|
path
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpClient for RemoteMcpClient {
|
||||||
|
async fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> {
|
||||||
|
// TODO: Implement remote call
|
||||||
|
Err(Error::NotImplemented(
|
||||||
|
"Remote MCP client is not implemented".to_string(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tool(&self, _call: McpToolCall) -> Result<McpToolResponse> {
|
||||||
|
// TODO: Implement remote call
|
||||||
|
Err(Error::NotImplemented(
|
||||||
|
"Remote MCP client is not implemented".to_string(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
87
crates/owlen-core/src/mcp/factory.rs
Normal file
87
crates/owlen-core/src/mcp/factory.rs
Normal file
@@ -0,0 +1,87 @@
|
|||||||
|
/// MCP Client Factory
|
||||||
|
///
|
||||||
|
/// Provides a unified interface for creating MCP clients based on configuration.
|
||||||
|
/// Supports switching between local (in-process) and remote (STDIO) execution modes.
|
||||||
|
use super::client::McpClient;
|
||||||
|
use super::{remote_client::RemoteMcpClient, LocalMcpClient};
|
||||||
|
use crate::config::{Config, McpMode};
|
||||||
|
use crate::tools::registry::ToolRegistry;
|
||||||
|
use crate::validation::SchemaValidator;
|
||||||
|
use crate::Result;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Factory for creating MCP clients based on configuration
|
||||||
|
pub struct McpClientFactory {
|
||||||
|
config: Arc<Config>,
|
||||||
|
registry: Arc<ToolRegistry>,
|
||||||
|
validator: Arc<SchemaValidator>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl McpClientFactory {
|
||||||
|
pub fn new(
|
||||||
|
config: Arc<Config>,
|
||||||
|
registry: Arc<ToolRegistry>,
|
||||||
|
validator: Arc<SchemaValidator>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
config,
|
||||||
|
registry,
|
||||||
|
validator,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create an MCP client based on the current configuration
|
||||||
|
pub fn create(&self) -> Result<Box<dyn McpClient>> {
|
||||||
|
match self.config.mcp.mode {
|
||||||
|
McpMode::Legacy => {
|
||||||
|
// Use local in-process client
|
||||||
|
Ok(Box::new(LocalMcpClient::new(
|
||||||
|
self.registry.clone(),
|
||||||
|
self.validator.clone(),
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
McpMode::Enabled => {
|
||||||
|
// Attempt to use remote client, fall back to local if unavailable
|
||||||
|
match RemoteMcpClient::new() {
|
||||||
|
Ok(client) => Ok(Box::new(client)),
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("Warning: Failed to start remote MCP client: {}. Falling back to local mode.", e);
|
||||||
|
Ok(Box::new(LocalMcpClient::new(
|
||||||
|
self.registry.clone(),
|
||||||
|
self.validator.clone(),
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if remote MCP mode is available
|
||||||
|
pub fn is_remote_available() -> bool {
|
||||||
|
RemoteMcpClient::new().is_ok()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_factory_creates_local_client_in_legacy_mode() {
|
||||||
|
let mut config = Config::default();
|
||||||
|
config.mcp.mode = McpMode::Legacy;
|
||||||
|
|
||||||
|
let ui = Arc::new(crate::ui::NoOpUiController);
|
||||||
|
let registry = Arc::new(ToolRegistry::new(
|
||||||
|
Arc::new(tokio::sync::Mutex::new(config.clone())),
|
||||||
|
ui,
|
||||||
|
));
|
||||||
|
let validator = Arc::new(SchemaValidator::new());
|
||||||
|
|
||||||
|
let factory = McpClientFactory::new(Arc::new(config), registry, validator);
|
||||||
|
|
||||||
|
// Should create without error
|
||||||
|
let result = factory.create();
|
||||||
|
assert!(result.is_ok());
|
||||||
|
}
|
||||||
|
}
|
||||||
217
crates/owlen-core/src/mcp/permission.rs
Normal file
217
crates/owlen-core/src/mcp/permission.rs
Normal file
@@ -0,0 +1,217 @@
|
|||||||
|
/// Permission and Safety Layer for MCP
|
||||||
|
///
|
||||||
|
/// This module provides runtime enforcement of security policies for tool execution.
|
||||||
|
/// It wraps MCP clients to filter/whitelist tool calls, log invocations, and prompt for consent.
|
||||||
|
use super::client::McpClient;
|
||||||
|
use super::{McpToolCall, McpToolDescriptor, McpToolResponse};
|
||||||
|
use crate::config::Config;
|
||||||
|
use crate::{Error, Result};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use std::collections::HashSet;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Callback for requesting user consent for dangerous operations
|
||||||
|
pub type ConsentCallback = Arc<dyn Fn(&str, &McpToolCall) -> bool + Send + Sync>;
|
||||||
|
|
||||||
|
/// Callback for logging tool invocations
|
||||||
|
pub type LogCallback = Arc<dyn Fn(&str, &McpToolCall, &Result<McpToolResponse>) + Send + Sync>;
|
||||||
|
|
||||||
|
/// Permission-enforcing wrapper around an MCP client
|
||||||
|
pub struct PermissionLayer {
|
||||||
|
inner: Box<dyn McpClient>,
|
||||||
|
config: Arc<Config>,
|
||||||
|
consent_callback: Option<ConsentCallback>,
|
||||||
|
log_callback: Option<LogCallback>,
|
||||||
|
allowed_tools: HashSet<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PermissionLayer {
|
||||||
|
/// Create a new permission layer wrapping the given client
|
||||||
|
pub fn new(inner: Box<dyn McpClient>, config: Arc<Config>) -> Self {
|
||||||
|
let allowed_tools = config.security.allowed_tools.iter().cloned().collect();
|
||||||
|
|
||||||
|
Self {
|
||||||
|
inner,
|
||||||
|
config,
|
||||||
|
consent_callback: None,
|
||||||
|
log_callback: None,
|
||||||
|
allowed_tools,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Set a callback for requesting user consent
|
||||||
|
pub fn with_consent_callback(mut self, callback: ConsentCallback) -> Self {
|
||||||
|
self.consent_callback = Some(callback);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Set a callback for logging tool invocations
|
||||||
|
pub fn with_log_callback(mut self, callback: LogCallback) -> Self {
|
||||||
|
self.log_callback = Some(callback);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if a tool requires dangerous filesystem operations
|
||||||
|
fn requires_dangerous_filesystem(&self, tool_name: &str) -> bool {
|
||||||
|
matches!(
|
||||||
|
tool_name,
|
||||||
|
"resources/write" | "resources/delete" | "file_write" | "file_delete"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if a tool is allowed by security policy
|
||||||
|
fn is_tool_allowed(&self, tool_descriptor: &McpToolDescriptor) -> bool {
|
||||||
|
// Check if tool requires filesystem access
|
||||||
|
for fs_perm in &tool_descriptor.requires_filesystem {
|
||||||
|
if !self.allowed_tools.contains(fs_perm) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if tool requires network access
|
||||||
|
if tool_descriptor.requires_network && !self.allowed_tools.contains("web_search") {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Request user consent for a tool call
|
||||||
|
fn request_consent(&self, tool_name: &str, call: &McpToolCall) -> bool {
|
||||||
|
if let Some(ref callback) = self.consent_callback {
|
||||||
|
callback(tool_name, call)
|
||||||
|
} else {
|
||||||
|
// If no callback is set, deny dangerous operations by default
|
||||||
|
!self.requires_dangerous_filesystem(tool_name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Log a tool invocation
|
||||||
|
fn log_invocation(
|
||||||
|
&self,
|
||||||
|
tool_name: &str,
|
||||||
|
call: &McpToolCall,
|
||||||
|
result: &Result<McpToolResponse>,
|
||||||
|
) {
|
||||||
|
if let Some(ref callback) = self.log_callback {
|
||||||
|
callback(tool_name, call, result);
|
||||||
|
} else {
|
||||||
|
// Default logging to stderr
|
||||||
|
match result {
|
||||||
|
Ok(resp) => {
|
||||||
|
eprintln!(
|
||||||
|
"[MCP] Tool '{}' executed successfully ({}ms)",
|
||||||
|
tool_name, resp.duration_ms
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("[MCP] Tool '{}' failed: {}", tool_name, e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl McpClient for PermissionLayer {
|
||||||
|
async fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> {
|
||||||
|
let tools = self.inner.list_tools().await?;
|
||||||
|
// Filter tools based on security policy
|
||||||
|
Ok(tools
|
||||||
|
.into_iter()
|
||||||
|
.filter(|tool| self.is_tool_allowed(tool))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tool(&self, call: McpToolCall) -> Result<McpToolResponse> {
|
||||||
|
// Check if tool requires consent
|
||||||
|
if self.requires_dangerous_filesystem(&call.name)
|
||||||
|
&& self.config.privacy.require_consent_per_session
|
||||||
|
&& !self.request_consent(&call.name, &call)
|
||||||
|
{
|
||||||
|
let result = Err(Error::PermissionDenied(format!(
|
||||||
|
"User denied consent for tool '{}'",
|
||||||
|
call.name
|
||||||
|
)));
|
||||||
|
self.log_invocation(&call.name, &call, &result);
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute the tool call
|
||||||
|
let result = self.inner.call_tool(call.clone()).await;
|
||||||
|
|
||||||
|
// Log the invocation
|
||||||
|
self.log_invocation(&call.name, &call, &result);
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::mcp::LocalMcpClient;
|
||||||
|
use crate::tools::registry::ToolRegistry;
|
||||||
|
use crate::validation::SchemaValidator;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_permission_layer_filters_dangerous_tools() {
|
||||||
|
let config = Arc::new(Config::default());
|
||||||
|
let ui = Arc::new(crate::ui::NoOpUiController);
|
||||||
|
let registry = Arc::new(ToolRegistry::new(
|
||||||
|
Arc::new(tokio::sync::Mutex::new((*config).clone())),
|
||||||
|
ui,
|
||||||
|
));
|
||||||
|
let validator = Arc::new(SchemaValidator::new());
|
||||||
|
let client = Box::new(LocalMcpClient::new(registry, validator));
|
||||||
|
|
||||||
|
let mut config_mut = (*config).clone();
|
||||||
|
// Disallow file operations
|
||||||
|
config_mut.security.allowed_tools = vec!["web_search".to_string()];
|
||||||
|
|
||||||
|
let permission_layer = PermissionLayer::new(client, Arc::new(config_mut));
|
||||||
|
|
||||||
|
let tools = permission_layer.list_tools().await.unwrap();
|
||||||
|
|
||||||
|
// Should not include file_write or file_delete tools
|
||||||
|
assert!(!tools.iter().any(|t| t.name.contains("write")));
|
||||||
|
assert!(!tools.iter().any(|t| t.name.contains("delete")));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_consent_callback_is_invoked() {
|
||||||
|
let config = Arc::new(Config::default());
|
||||||
|
let ui = Arc::new(crate::ui::NoOpUiController);
|
||||||
|
let registry = Arc::new(ToolRegistry::new(
|
||||||
|
Arc::new(tokio::sync::Mutex::new((*config).clone())),
|
||||||
|
ui,
|
||||||
|
));
|
||||||
|
let validator = Arc::new(SchemaValidator::new());
|
||||||
|
let client = Box::new(LocalMcpClient::new(registry, validator));
|
||||||
|
|
||||||
|
let consent_called = Arc::new(AtomicBool::new(false));
|
||||||
|
let consent_called_clone = consent_called.clone();
|
||||||
|
|
||||||
|
let consent_callback: ConsentCallback = Arc::new(move |_tool, _call| {
|
||||||
|
consent_called_clone.store(true, Ordering::SeqCst);
|
||||||
|
false // Deny
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut config_mut = (*config).clone();
|
||||||
|
config_mut.privacy.require_consent_per_session = true;
|
||||||
|
|
||||||
|
let permission_layer = PermissionLayer::new(client, Arc::new(config_mut))
|
||||||
|
.with_consent_callback(consent_callback);
|
||||||
|
|
||||||
|
let call = McpToolCall {
|
||||||
|
name: "resources/write".to_string(),
|
||||||
|
arguments: serde_json::json!({"path": "test.txt", "content": "hello"}),
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = permission_layer.call_tool(call).await;
|
||||||
|
|
||||||
|
assert!(consent_called.load(Ordering::SeqCst));
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
}
|
||||||
389
crates/owlen-core/src/mcp/protocol.rs
Normal file
389
crates/owlen-core/src/mcp/protocol.rs
Normal file
@@ -0,0 +1,389 @@
|
|||||||
|
/// MCP Protocol Definitions
|
||||||
|
///
|
||||||
|
/// This module defines the JSON-RPC protocol contracts for the Model Context Protocol (MCP).
|
||||||
|
/// It includes request/response schemas, error codes, and versioning semantics.
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// MCP Protocol version - uses semantic versioning
|
||||||
|
pub const PROTOCOL_VERSION: &str = "1.0.0";
|
||||||
|
|
||||||
|
/// JSON-RPC version constant
|
||||||
|
pub const JSONRPC_VERSION: &str = "2.0";
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Error Codes and Handling
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// Standard JSON-RPC error codes following the spec
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
pub struct ErrorCode(pub i64);
|
||||||
|
|
||||||
|
impl ErrorCode {
|
||||||
|
// Standard JSON-RPC 2.0 errors
|
||||||
|
pub const PARSE_ERROR: Self = Self(-32700);
|
||||||
|
pub const INVALID_REQUEST: Self = Self(-32600);
|
||||||
|
pub const METHOD_NOT_FOUND: Self = Self(-32601);
|
||||||
|
pub const INVALID_PARAMS: Self = Self(-32602);
|
||||||
|
pub const INTERNAL_ERROR: Self = Self(-32603);
|
||||||
|
|
||||||
|
// MCP-specific errors (range -32000 to -32099)
|
||||||
|
pub const TOOL_NOT_FOUND: Self = Self(-32000);
|
||||||
|
pub const TOOL_EXECUTION_FAILED: Self = Self(-32001);
|
||||||
|
pub const PERMISSION_DENIED: Self = Self(-32002);
|
||||||
|
pub const RESOURCE_NOT_FOUND: Self = Self(-32003);
|
||||||
|
pub const TIMEOUT: Self = Self(-32004);
|
||||||
|
pub const VALIDATION_ERROR: Self = Self(-32005);
|
||||||
|
pub const PATH_TRAVERSAL: Self = Self(-32006);
|
||||||
|
pub const RATE_LIMIT_EXCEEDED: Self = Self(-32007);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Structured error response
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RpcError {
|
||||||
|
pub code: i64,
|
||||||
|
pub message: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub data: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RpcError {
|
||||||
|
pub fn new(code: ErrorCode, message: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
code: code.0,
|
||||||
|
message: message.into(),
|
||||||
|
data: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_data(mut self, data: Value) -> Self {
|
||||||
|
self.data = Some(data);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn parse_error(message: impl Into<String>) -> Self {
|
||||||
|
Self::new(ErrorCode::PARSE_ERROR, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn invalid_request(message: impl Into<String>) -> Self {
|
||||||
|
Self::new(ErrorCode::INVALID_REQUEST, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn method_not_found(method: &str) -> Self {
|
||||||
|
Self::new(
|
||||||
|
ErrorCode::METHOD_NOT_FOUND,
|
||||||
|
format!("Method not found: {}", method),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn invalid_params(message: impl Into<String>) -> Self {
|
||||||
|
Self::new(ErrorCode::INVALID_PARAMS, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn internal_error(message: impl Into<String>) -> Self {
|
||||||
|
Self::new(ErrorCode::INTERNAL_ERROR, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn tool_not_found(tool_name: &str) -> Self {
|
||||||
|
Self::new(
|
||||||
|
ErrorCode::TOOL_NOT_FOUND,
|
||||||
|
format!("Tool not found: {}", tool_name),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn permission_denied(message: impl Into<String>) -> Self {
|
||||||
|
Self::new(ErrorCode::PERMISSION_DENIED, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn path_traversal() -> Self {
|
||||||
|
Self::new(ErrorCode::PATH_TRAVERSAL, "Path traversal attempt detected")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Request/Response Structures
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// JSON-RPC request structure
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RpcRequest {
|
||||||
|
pub jsonrpc: String,
|
||||||
|
pub id: RequestId,
|
||||||
|
pub method: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub params: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RpcRequest {
|
||||||
|
pub fn new(id: RequestId, method: impl Into<String>, params: Option<Value>) -> Self {
|
||||||
|
Self {
|
||||||
|
jsonrpc: JSONRPC_VERSION.to_string(),
|
||||||
|
id,
|
||||||
|
method: method.into(),
|
||||||
|
params,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// JSON-RPC response structure (success)
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RpcResponse {
|
||||||
|
pub jsonrpc: String,
|
||||||
|
pub id: RequestId,
|
||||||
|
pub result: Value,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RpcResponse {
|
||||||
|
pub fn new(id: RequestId, result: Value) -> Self {
|
||||||
|
Self {
|
||||||
|
jsonrpc: JSONRPC_VERSION.to_string(),
|
||||||
|
id,
|
||||||
|
result,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// JSON-RPC error response
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RpcErrorResponse {
|
||||||
|
pub jsonrpc: String,
|
||||||
|
pub id: RequestId,
|
||||||
|
pub error: RpcError,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RpcErrorResponse {
|
||||||
|
pub fn new(id: RequestId, error: RpcError) -> Self {
|
||||||
|
Self {
|
||||||
|
jsonrpc: JSONRPC_VERSION.to_string(),
|
||||||
|
id,
|
||||||
|
error,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// JSON‑RPC notification (no id). Used for streaming partial results.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RpcNotification {
|
||||||
|
pub jsonrpc: String,
|
||||||
|
pub method: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub params: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RpcNotification {
|
||||||
|
pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
|
||||||
|
Self {
|
||||||
|
jsonrpc: JSONRPC_VERSION.to_string(),
|
||||||
|
method: method.into(),
|
||||||
|
params,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Request ID can be string, number, or null
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
||||||
|
#[serde(untagged)]
|
||||||
|
pub enum RequestId {
|
||||||
|
Number(u64),
|
||||||
|
String(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<u64> for RequestId {
|
||||||
|
fn from(n: u64) -> Self {
|
||||||
|
Self::Number(n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<String> for RequestId {
|
||||||
|
fn from(s: String) -> Self {
|
||||||
|
Self::String(s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// MCP Method Names
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// Standard MCP methods
|
||||||
|
pub mod methods {
|
||||||
|
pub const INITIALIZE: &str = "initialize";
|
||||||
|
pub const TOOLS_LIST: &str = "tools/list";
|
||||||
|
pub const TOOLS_CALL: &str = "tools/call";
|
||||||
|
pub const RESOURCES_LIST: &str = "resources/list";
|
||||||
|
pub const RESOURCES_GET: &str = "resources/get";
|
||||||
|
pub const RESOURCES_WRITE: &str = "resources/write";
|
||||||
|
pub const RESOURCES_DELETE: &str = "resources/delete";
|
||||||
|
pub const MODELS_LIST: &str = "models/list";
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Initialization Protocol
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// Initialize request parameters
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct InitializeParams {
|
||||||
|
pub protocol_version: String,
|
||||||
|
pub client_info: ClientInfo,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub capabilities: Option<ClientCapabilities>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for InitializeParams {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
protocol_version: PROTOCOL_VERSION.to_string(),
|
||||||
|
client_info: ClientInfo {
|
||||||
|
name: "owlen".to_string(),
|
||||||
|
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||||
|
},
|
||||||
|
capabilities: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Client information
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ClientInfo {
|
||||||
|
pub name: String,
|
||||||
|
pub version: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Client capabilities
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct ClientCapabilities {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub supports_streaming: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub supports_cancellation: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Initialize response
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct InitializeResult {
|
||||||
|
pub protocol_version: String,
|
||||||
|
pub server_info: ServerInfo,
|
||||||
|
pub capabilities: ServerCapabilities,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Server information
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ServerInfo {
|
||||||
|
pub name: String,
|
||||||
|
pub version: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Server capabilities
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct ServerCapabilities {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub supports_tools: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub supports_resources: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub supports_streaming: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Tool Call Protocol
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// Parameters for tools/list
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct ToolsListParams {
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub filter: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameters for tools/call
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolsCallParams {
|
||||||
|
pub name: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub arguments: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of tools/call
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolsCallResult {
|
||||||
|
pub success: bool,
|
||||||
|
pub output: Value,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub error: Option<String>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub metadata: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Resource Protocol
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// Parameters for resources/list
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ResourcesListParams {
|
||||||
|
pub path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameters for resources/get
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ResourcesGetParams {
|
||||||
|
pub path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameters for resources/write
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ResourcesWriteParams {
|
||||||
|
pub path: String,
|
||||||
|
pub content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parameters for resources/delete
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ResourcesDeleteParams {
|
||||||
|
pub path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ============================================================================
|
||||||
|
// Versioning and Compatibility
|
||||||
|
// ============================================================================
|
||||||
|
|
||||||
|
/// Check if a protocol version is compatible
|
||||||
|
pub fn is_compatible(client_version: &str, server_version: &str) -> bool {
|
||||||
|
// For now, simple exact match on major version
|
||||||
|
let client_major = client_version.split('.').next().unwrap_or("0");
|
||||||
|
let server_major = server_version.split('.').next().unwrap_or("0");
|
||||||
|
client_major == server_major
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_error_codes() {
|
||||||
|
let err = RpcError::tool_not_found("test_tool");
|
||||||
|
assert_eq!(err.code, ErrorCode::TOOL_NOT_FOUND.0);
|
||||||
|
assert!(err.message.contains("test_tool"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_version_compatibility() {
|
||||||
|
assert!(is_compatible("1.0.0", "1.0.0"));
|
||||||
|
assert!(is_compatible("1.0.0", "1.1.0"));
|
||||||
|
assert!(is_compatible("1.2.5", "1.0.0"));
|
||||||
|
assert!(!is_compatible("1.0.0", "2.0.0"));
|
||||||
|
assert!(!is_compatible("2.0.0", "1.0.0"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_request_serialization() {
|
||||||
|
let req = RpcRequest::new(
|
||||||
|
RequestId::Number(1),
|
||||||
|
"tools/call",
|
||||||
|
Some(serde_json::json!({"name": "test"})),
|
||||||
|
);
|
||||||
|
let json = serde_json::to_string(&req).unwrap();
|
||||||
|
assert!(json.contains("\"jsonrpc\":\"2.0\""));
|
||||||
|
assert!(json.contains("\"method\":\"tools/call\""));
|
||||||
|
}
|
||||||
|
}
|
||||||
254
crates/owlen-core/src/mcp/remote_client.rs
Normal file
254
crates/owlen-core/src/mcp/remote_client.rs
Normal file
@@ -0,0 +1,254 @@
|
|||||||
|
use super::protocol::methods;
|
||||||
|
use super::protocol::{RequestId, RpcErrorResponse, RpcRequest, RpcResponse, PROTOCOL_VERSION};
|
||||||
|
use super::{McpClient, McpToolCall, McpToolDescriptor, McpToolResponse};
|
||||||
|
use crate::types::ModelInfo;
|
||||||
|
use crate::{Error, Provider, Result};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||||
|
use tokio::process::{Child, Command};
|
||||||
|
use tokio::sync::Mutex;
|
||||||
|
// Provider trait is already imported via the earlier use statement.
|
||||||
|
use crate::types::{ChatResponse, Message, Role};
|
||||||
|
use futures::stream;
|
||||||
|
use futures::StreamExt;
|
||||||
|
|
||||||
|
/// Client that talks to the external `owlen-mcp-server` over STDIO.
|
||||||
|
pub struct RemoteMcpClient {
|
||||||
|
// Child process handling the server (kept alive for the duration of the client).
|
||||||
|
#[allow(dead_code)]
|
||||||
|
child: Arc<Mutex<Child>>, // guarded for mutable access across calls
|
||||||
|
// Writer to server stdin.
|
||||||
|
stdin: Arc<Mutex<tokio::process::ChildStdin>>, // async write
|
||||||
|
// Reader for server stdout.
|
||||||
|
stdout: Arc<Mutex<BufReader<tokio::process::ChildStdout>>>,
|
||||||
|
// Incrementing request identifier.
|
||||||
|
next_id: AtomicU64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RemoteMcpClient {
|
||||||
|
/// Spawn the MCP server binary and prepare communication channels.
|
||||||
|
pub fn new() -> Result<Self> {
|
||||||
|
// Locate the binary – it is built by Cargo into target/debug.
|
||||||
|
// The test binary runs inside the crate directory, so we check a couple of relative locations.
|
||||||
|
// Attempt to locate the server binary; if unavailable we will fall back to launching via `cargo run`.
|
||||||
|
let _ = ();
|
||||||
|
// Resolve absolute path based on workspace root to avoid cwd dependence.
|
||||||
|
// The MCP server binary lives in the workspace's `target/debug` directory.
|
||||||
|
// Historically the binary was named `owlen-mcp-server`, but it has been
|
||||||
|
// renamed to `owlen-mcp-llm-server`. We attempt to locate the new name
|
||||||
|
// first and fall back to the legacy name for compatibility.
|
||||||
|
let workspace_root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||||
|
.join("../..")
|
||||||
|
.canonicalize()
|
||||||
|
.map_err(Error::Io)?;
|
||||||
|
let candidates = [
|
||||||
|
"target/debug/owlen-mcp-llm-server",
|
||||||
|
"target/debug/owlen-mcp-server",
|
||||||
|
];
|
||||||
|
let mut binary_path = None;
|
||||||
|
for rel in &candidates {
|
||||||
|
let p = workspace_root.join(rel);
|
||||||
|
if p.exists() {
|
||||||
|
binary_path = Some(p);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let binary_path = binary_path.ok_or_else(|| {
|
||||||
|
Error::NotImplemented(format!(
|
||||||
|
"owlen-mcp server binary not found; checked {} and {}",
|
||||||
|
candidates[0], candidates[1]
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
if !binary_path.exists() {
|
||||||
|
return Err(Error::NotImplemented(format!(
|
||||||
|
"owlen-mcp-server binary not found at {}",
|
||||||
|
binary_path.display()
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
// Launch the already‑built server binary directly.
|
||||||
|
let mut child = Command::new(&binary_path)
|
||||||
|
.stdin(std::process::Stdio::piped())
|
||||||
|
.stdout(std::process::Stdio::piped())
|
||||||
|
.stderr(std::process::Stdio::inherit())
|
||||||
|
.spawn()
|
||||||
|
.map_err(Error::Io)?;
|
||||||
|
|
||||||
|
let stdin = child.stdin.take().ok_or_else(|| {
|
||||||
|
Error::Io(std::io::Error::other(
|
||||||
|
"Failed to capture stdin of MCP server",
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
let stdout = child.stdout.take().ok_or_else(|| {
|
||||||
|
Error::Io(std::io::Error::other(
|
||||||
|
"Failed to capture stdout of MCP server",
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
child: Arc::new(Mutex::new(child)),
|
||||||
|
stdin: Arc::new(Mutex::new(stdin)),
|
||||||
|
stdout: Arc::new(Mutex::new(BufReader::new(stdout))),
|
||||||
|
next_id: AtomicU64::new(1),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn send_rpc(&self, method: &str, params: serde_json::Value) -> Result<serde_json::Value> {
|
||||||
|
let id = RequestId::Number(self.next_id.fetch_add(1, Ordering::Relaxed));
|
||||||
|
let request = RpcRequest::new(id.clone(), method, Some(params));
|
||||||
|
let req_str = serde_json::to_string(&request)? + "\n";
|
||||||
|
{
|
||||||
|
let mut stdin = self.stdin.lock().await;
|
||||||
|
stdin.write_all(req_str.as_bytes()).await?;
|
||||||
|
stdin.flush().await?;
|
||||||
|
}
|
||||||
|
// Read a single line response
|
||||||
|
let mut line = String::new();
|
||||||
|
{
|
||||||
|
let mut stdout = self.stdout.lock().await;
|
||||||
|
stdout.read_line(&mut line).await?;
|
||||||
|
}
|
||||||
|
// Try to parse successful response first
|
||||||
|
if let Ok(resp) = serde_json::from_str::<RpcResponse>(&line) {
|
||||||
|
if resp.id == id {
|
||||||
|
return Ok(resp.result);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Fallback to error response
|
||||||
|
let err_resp: RpcErrorResponse =
|
||||||
|
serde_json::from_str(&line).map_err(Error::Serialization)?;
|
||||||
|
Err(Error::Network(format!(
|
||||||
|
"MCP server error {}: {}",
|
||||||
|
err_resp.error.code, err_resp.error.message
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl McpClient for RemoteMcpClient {
|
||||||
|
async fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> {
|
||||||
|
// Query the remote MCP server for its tool descriptors using the standard
|
||||||
|
// `tools/list` RPC method. The server returns a JSON array of
|
||||||
|
// `McpToolDescriptor` objects.
|
||||||
|
let result = self.send_rpc(methods::TOOLS_LIST, json!(null)).await?;
|
||||||
|
let descriptors: Vec<McpToolDescriptor> = serde_json::from_value(result)?;
|
||||||
|
Ok(descriptors)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tool(&self, call: McpToolCall) -> Result<McpToolResponse> {
|
||||||
|
// Local handling for simple resource tools to avoid needing the MCP server
|
||||||
|
// to implement them.
|
||||||
|
if call.name.starts_with("resources/get") {
|
||||||
|
let path = call
|
||||||
|
.arguments
|
||||||
|
.get("path")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("");
|
||||||
|
let content = std::fs::read_to_string(path).map_err(Error::Io)?;
|
||||||
|
return Ok(McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: true,
|
||||||
|
output: serde_json::json!(content),
|
||||||
|
metadata: std::collections::HashMap::new(),
|
||||||
|
duration_ms: 0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if call.name.starts_with("resources/list") {
|
||||||
|
let path = call
|
||||||
|
.arguments
|
||||||
|
.get("path")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or(".");
|
||||||
|
let mut names = Vec::new();
|
||||||
|
for entry in std::fs::read_dir(path).map_err(Error::Io)?.flatten() {
|
||||||
|
if let Some(name) = entry.file_name().to_str() {
|
||||||
|
names.push(name.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return Ok(McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: true,
|
||||||
|
output: serde_json::json!(names),
|
||||||
|
metadata: std::collections::HashMap::new(),
|
||||||
|
duration_ms: 0,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
// MCP server expects a generic "tools/call" method with a payload containing the
|
||||||
|
// specific tool name and its arguments. Wrap the incoming call accordingly.
|
||||||
|
let payload = serde_json::to_value(&call)?;
|
||||||
|
let result = self.send_rpc(methods::TOOLS_CALL, payload).await?;
|
||||||
|
// The server returns the tool's output directly; construct a matching response.
|
||||||
|
Ok(McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: true,
|
||||||
|
output: result,
|
||||||
|
metadata: std::collections::HashMap::new(),
|
||||||
|
duration_ms: 0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Provider implementation – forwards chat requests to the generate_text tool.
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Provider for RemoteMcpClient {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"mcp-llm-server"
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_models(&self) -> Result<Vec<ModelInfo>> {
|
||||||
|
let result = self.send_rpc(methods::MODELS_LIST, json!(null)).await?;
|
||||||
|
let models: Vec<ModelInfo> = serde_json::from_value(result)?;
|
||||||
|
Ok(models)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn chat(&self, request: crate::types::ChatRequest) -> Result<ChatResponse> {
|
||||||
|
// Use the streaming implementation and take the first response.
|
||||||
|
let mut stream = self.chat_stream(request).await?;
|
||||||
|
match stream.next().await {
|
||||||
|
Some(Ok(resp)) => Ok(resp),
|
||||||
|
Some(Err(e)) => Err(e),
|
||||||
|
None => Err(Error::Provider(anyhow::anyhow!("Empty chat stream"))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn chat_stream(
|
||||||
|
&self,
|
||||||
|
request: crate::types::ChatRequest,
|
||||||
|
) -> Result<crate::provider::ChatStream> {
|
||||||
|
// Build arguments matching the generate_text schema.
|
||||||
|
let args = serde_json::json!({
|
||||||
|
"messages": request.messages,
|
||||||
|
"temperature": request.parameters.temperature,
|
||||||
|
"max_tokens": request.parameters.max_tokens,
|
||||||
|
"model": request.model,
|
||||||
|
"stream": request.parameters.stream,
|
||||||
|
});
|
||||||
|
let call = McpToolCall {
|
||||||
|
name: "generate_text".to_string(),
|
||||||
|
arguments: args,
|
||||||
|
};
|
||||||
|
let resp = self.call_tool(call).await?;
|
||||||
|
// Build a ChatResponse from the tool output (assumed to be a string).
|
||||||
|
let content = resp.output.as_str().unwrap_or("").to_string();
|
||||||
|
let message = Message::new(Role::Assistant, content);
|
||||||
|
let chat_resp = ChatResponse {
|
||||||
|
message,
|
||||||
|
usage: None,
|
||||||
|
is_streaming: false,
|
||||||
|
is_final: true,
|
||||||
|
};
|
||||||
|
let stream = stream::once(async move { Ok(chat_resp) });
|
||||||
|
Ok(Box::pin(stream))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn health_check(&self) -> Result<()> {
|
||||||
|
// Simple ping using initialize method.
|
||||||
|
let params = serde_json::json!({"protocol_version": PROTOCOL_VERSION});
|
||||||
|
self.send_rpc("initialize", params).await.map(|_| ())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -17,7 +17,7 @@ pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatResponse>> + Send>>;
|
|||||||
/// use std::sync::Arc;
|
/// use std::sync::Arc;
|
||||||
/// use futures::Stream;
|
/// use futures::Stream;
|
||||||
/// use owlen_core::provider::{Provider, ProviderRegistry, ChatStream};
|
/// use owlen_core::provider::{Provider, ProviderRegistry, ChatStream};
|
||||||
/// use owlen_core::types::{ChatRequest, ChatResponse, ModelInfo, Message};
|
/// use owlen_core::types::{ChatRequest, ChatResponse, ModelInfo, Message, Role, ChatParameters};
|
||||||
/// use owlen_core::Result;
|
/// use owlen_core::Result;
|
||||||
///
|
///
|
||||||
/// // 1. Create a mock provider
|
/// // 1. Create a mock provider
|
||||||
@@ -31,18 +31,23 @@ pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatResponse>> + Send>>;
|
|||||||
///
|
///
|
||||||
/// async fn list_models(&self) -> Result<Vec<ModelInfo>> {
|
/// async fn list_models(&self) -> Result<Vec<ModelInfo>> {
|
||||||
/// Ok(vec![ModelInfo {
|
/// Ok(vec![ModelInfo {
|
||||||
|
/// id: "mock-model".to_string(),
|
||||||
/// provider: "mock".to_string(),
|
/// provider: "mock".to_string(),
|
||||||
/// name: "mock-model".to_string(),
|
/// name: "mock-model".to_string(),
|
||||||
/// ..Default::default()
|
/// description: None,
|
||||||
|
/// context_window: None,
|
||||||
|
/// capabilities: vec![],
|
||||||
|
/// supports_tools: false,
|
||||||
/// }])
|
/// }])
|
||||||
/// }
|
/// }
|
||||||
///
|
///
|
||||||
/// async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
|
/// async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
|
||||||
/// let content = format!("Response to: {}", request.messages.last().unwrap().content);
|
/// let content = format!("Response to: {}", request.messages.last().unwrap().content);
|
||||||
/// Ok(ChatResponse {
|
/// Ok(ChatResponse {
|
||||||
/// model: request.model,
|
/// message: Message::new(Role::Assistant, content),
|
||||||
/// message: Message { role: "assistant".to_string(), content, ..Default::default() },
|
/// usage: None,
|
||||||
/// ..Default::default()
|
/// is_streaming: false,
|
||||||
|
/// is_final: true,
|
||||||
/// })
|
/// })
|
||||||
/// }
|
/// }
|
||||||
///
|
///
|
||||||
@@ -67,8 +72,9 @@ pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatResponse>> + Send>>;
|
|||||||
///
|
///
|
||||||
/// let request = ChatRequest {
|
/// let request = ChatRequest {
|
||||||
/// model: "mock-model".to_string(),
|
/// model: "mock-model".to_string(),
|
||||||
/// messages: vec![Message { role: "user".to_string(), content: "Hello".to_string(), ..Default::default() }],
|
/// messages: vec![Message::new(Role::User, "Hello".to_string())],
|
||||||
/// ..Default::default()
|
/// parameters: ChatParameters::default(),
|
||||||
|
/// tools: None,
|
||||||
/// };
|
/// };
|
||||||
///
|
///
|
||||||
/// let response = provider.chat(request).await.unwrap();
|
/// let response = provider.chat(request).await.unwrap();
|
||||||
|
|||||||
212
crates/owlen-core/src/sandbox.rs
Normal file
212
crates/owlen-core/src/sandbox.rs
Normal file
@@ -0,0 +1,212 @@
|
|||||||
|
use std::path::PathBuf;
|
||||||
|
use std::process::{Command, Stdio};
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use anyhow::{bail, Context, Result};
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
/// Configuration options for sandboxed process execution.
|
||||||
|
#[derive(Clone, Debug)]
|
||||||
|
pub struct SandboxConfig {
|
||||||
|
pub allow_network: bool,
|
||||||
|
pub allow_paths: Vec<PathBuf>,
|
||||||
|
pub readonly_paths: Vec<PathBuf>,
|
||||||
|
pub timeout_seconds: u64,
|
||||||
|
pub max_memory_mb: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for SandboxConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
allow_network: false,
|
||||||
|
allow_paths: Vec::new(),
|
||||||
|
readonly_paths: Vec::new(),
|
||||||
|
timeout_seconds: 30,
|
||||||
|
max_memory_mb: 512,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Wrapper around a bubblewrap sandbox instance.
|
||||||
|
///
|
||||||
|
/// Memory limits are enforced via:
|
||||||
|
/// - bwrap's --rlimit-as (version >= 0.12.0)
|
||||||
|
/// - prlimit wrapper (fallback for older bwrap versions)
|
||||||
|
/// - timeout mechanism (always enforced as last resort)
|
||||||
|
pub struct SandboxedProcess {
|
||||||
|
temp_dir: TempDir,
|
||||||
|
config: SandboxConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SandboxedProcess {
|
||||||
|
pub fn new(config: SandboxConfig) -> Result<Self> {
|
||||||
|
let temp_dir = TempDir::new().context("Failed to create temp directory")?;
|
||||||
|
|
||||||
|
which::which("bwrap")
|
||||||
|
.context("bubblewrap not found. Install with: sudo apt install bubblewrap")?;
|
||||||
|
|
||||||
|
Ok(Self { temp_dir, config })
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn execute(&self, command: &str, args: &[&str]) -> Result<SandboxResult> {
|
||||||
|
let supports_rlimit = self.supports_rlimit_as();
|
||||||
|
let use_prlimit = !supports_rlimit && which::which("prlimit").is_ok();
|
||||||
|
|
||||||
|
let mut cmd = if use_prlimit {
|
||||||
|
// Use prlimit wrapper for older bwrap versions
|
||||||
|
let mut prlimit_cmd = Command::new("prlimit");
|
||||||
|
let memory_limit_bytes = self
|
||||||
|
.config
|
||||||
|
.max_memory_mb
|
||||||
|
.saturating_mul(1024)
|
||||||
|
.saturating_mul(1024);
|
||||||
|
prlimit_cmd.arg(format!("--as={}", memory_limit_bytes));
|
||||||
|
prlimit_cmd.arg("bwrap");
|
||||||
|
prlimit_cmd
|
||||||
|
} else {
|
||||||
|
Command::new("bwrap")
|
||||||
|
};
|
||||||
|
|
||||||
|
cmd.args(["--unshare-all", "--die-with-parent", "--new-session"]);
|
||||||
|
|
||||||
|
if self.config.allow_network {
|
||||||
|
cmd.arg("--share-net");
|
||||||
|
} else {
|
||||||
|
cmd.arg("--unshare-net");
|
||||||
|
}
|
||||||
|
|
||||||
|
cmd.args(["--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp"]);
|
||||||
|
|
||||||
|
// Bind essential system paths readonly for executables and libraries
|
||||||
|
let system_paths = ["/usr", "/bin", "/lib", "/lib64", "/etc"];
|
||||||
|
for sys_path in &system_paths {
|
||||||
|
let path = std::path::Path::new(sys_path);
|
||||||
|
if path.exists() {
|
||||||
|
cmd.arg("--ro-bind").arg(sys_path).arg(sys_path);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Bind /run for DNS resolution (resolv.conf may be a symlink to /run/systemd/resolve/*)
|
||||||
|
if std::path::Path::new("/run").exists() {
|
||||||
|
cmd.arg("--ro-bind").arg("/run").arg("/run");
|
||||||
|
}
|
||||||
|
|
||||||
|
for path in &self.config.allow_paths {
|
||||||
|
let path_host = path.to_string_lossy().into_owned();
|
||||||
|
let path_guest = path_host.clone();
|
||||||
|
cmd.arg("--bind").arg(&path_host).arg(&path_guest);
|
||||||
|
}
|
||||||
|
|
||||||
|
for path in &self.config.readonly_paths {
|
||||||
|
let path_host = path.to_string_lossy().into_owned();
|
||||||
|
let path_guest = path_host.clone();
|
||||||
|
cmd.arg("--ro-bind").arg(&path_host).arg(&path_guest);
|
||||||
|
}
|
||||||
|
|
||||||
|
let work_dir = self.temp_dir.path().to_string_lossy().into_owned();
|
||||||
|
cmd.arg("--bind").arg(&work_dir).arg("/work");
|
||||||
|
cmd.arg("--chdir").arg("/work");
|
||||||
|
|
||||||
|
// Add memory limits via bwrap's --rlimit-as if supported (version >= 0.12.0)
|
||||||
|
// If not supported, we use prlimit wrapper (set earlier)
|
||||||
|
if supports_rlimit && !use_prlimit {
|
||||||
|
let memory_limit_bytes = self
|
||||||
|
.config
|
||||||
|
.max_memory_mb
|
||||||
|
.saturating_mul(1024)
|
||||||
|
.saturating_mul(1024);
|
||||||
|
let memory_soft = memory_limit_bytes.to_string();
|
||||||
|
let memory_hard = memory_limit_bytes.to_string();
|
||||||
|
cmd.arg("--rlimit-as").arg(&memory_soft).arg(&memory_hard);
|
||||||
|
}
|
||||||
|
|
||||||
|
cmd.arg(command);
|
||||||
|
cmd.args(args);
|
||||||
|
|
||||||
|
let start = Instant::now();
|
||||||
|
let timeout = Duration::from_secs(self.config.timeout_seconds);
|
||||||
|
|
||||||
|
// Spawn the process instead of waiting immediately
|
||||||
|
let mut child = cmd
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped())
|
||||||
|
.spawn()
|
||||||
|
.context("Failed to spawn sandboxed command")?;
|
||||||
|
|
||||||
|
let mut was_timeout = false;
|
||||||
|
|
||||||
|
// Wait for the child with timeout
|
||||||
|
let output = loop {
|
||||||
|
match child.try_wait() {
|
||||||
|
Ok(Some(_status)) => {
|
||||||
|
// Process exited
|
||||||
|
let output = child
|
||||||
|
.wait_with_output()
|
||||||
|
.context("Failed to collect process output")?;
|
||||||
|
break output;
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
// Process still running, check timeout
|
||||||
|
if start.elapsed() >= timeout {
|
||||||
|
// Timeout exceeded, kill the process
|
||||||
|
was_timeout = true;
|
||||||
|
child.kill().context("Failed to kill timed-out process")?;
|
||||||
|
// Wait for the killed process to exit
|
||||||
|
let output = child
|
||||||
|
.wait_with_output()
|
||||||
|
.context("Failed to collect output from killed process")?;
|
||||||
|
break output;
|
||||||
|
}
|
||||||
|
// Sleep briefly before checking again
|
||||||
|
std::thread::sleep(Duration::from_millis(50));
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
bail!("Failed to check process status: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let duration = start.elapsed();
|
||||||
|
|
||||||
|
Ok(SandboxResult {
|
||||||
|
stdout: String::from_utf8_lossy(&output.stdout).to_string(),
|
||||||
|
stderr: String::from_utf8_lossy(&output.stderr).to_string(),
|
||||||
|
exit_code: output.status.code().unwrap_or(-1),
|
||||||
|
duration,
|
||||||
|
was_timeout,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if bubblewrap supports --rlimit-as option (version >= 0.12.0)
|
||||||
|
fn supports_rlimit_as(&self) -> bool {
|
||||||
|
// Try to get bwrap version
|
||||||
|
let output = Command::new("bwrap").arg("--version").output();
|
||||||
|
|
||||||
|
if let Ok(output) = output {
|
||||||
|
let version_str = String::from_utf8_lossy(&output.stdout);
|
||||||
|
// Parse version like "bubblewrap 0.11.0" or "0.11.0"
|
||||||
|
if let Some(version_part) = version_str.split_whitespace().last() {
|
||||||
|
if let Some((major, rest)) = version_part.split_once('.') {
|
||||||
|
if let Some((minor, _patch)) = rest.split_once('.') {
|
||||||
|
if let (Ok(maj), Ok(min)) = (major.parse::<u32>(), minor.parse::<u32>()) {
|
||||||
|
// --rlimit-as was added in 0.12.0
|
||||||
|
return maj > 0 || (maj == 0 && min >= 12);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// If we can't determine the version, assume it doesn't support it (safer default)
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct SandboxResult {
|
||||||
|
pub stdout: String,
|
||||||
|
pub stderr: String,
|
||||||
|
pub exit_code: i32,
|
||||||
|
pub duration: Duration,
|
||||||
|
pub was_timeout: bool,
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -1,19 +1,26 @@
|
|||||||
//! Session persistence and storage management
|
//! Session persistence and storage management backed by SQLite
|
||||||
|
|
||||||
use crate::types::Conversation;
|
use crate::types::Conversation;
|
||||||
use crate::{Error, Result};
|
use crate::{Error, Result};
|
||||||
|
use aes_gcm::aead::{Aead, KeyInit};
|
||||||
|
use aes_gcm::{Aes256Gcm, Nonce};
|
||||||
|
use ring::rand::{SecureRandom, SystemRandom};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
|
||||||
|
use sqlx::{Pool, Row, Sqlite};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
|
use std::io::IsTerminal;
|
||||||
|
use std::io::{self, Write};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::time::SystemTime;
|
use std::str::FromStr;
|
||||||
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
/// Metadata about a saved session
|
/// Metadata about a saved session
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct SessionMeta {
|
pub struct SessionMeta {
|
||||||
/// Session file path
|
|
||||||
pub path: PathBuf,
|
|
||||||
/// Conversation ID
|
/// Conversation ID
|
||||||
pub id: uuid::Uuid,
|
pub id: Uuid,
|
||||||
/// Optional session name
|
/// Optional session name
|
||||||
pub name: Option<String>,
|
pub name: Option<String>,
|
||||||
/// Optional AI-generated description
|
/// Optional AI-generated description
|
||||||
@@ -28,282 +35,525 @@ pub struct SessionMeta {
|
|||||||
pub updated_at: SystemTime,
|
pub updated_at: SystemTime,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Storage manager for persisting conversations
|
/// Storage manager for persisting conversations in SQLite
|
||||||
pub struct StorageManager {
|
pub struct StorageManager {
|
||||||
sessions_dir: PathBuf,
|
pool: Pool<Sqlite>,
|
||||||
|
database_path: PathBuf,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl StorageManager {
|
impl StorageManager {
|
||||||
/// Create a new storage manager with the default sessions directory
|
/// Create a new storage manager using the default database path
|
||||||
pub fn new() -> Result<Self> {
|
pub async fn new() -> Result<Self> {
|
||||||
let sessions_dir = Self::default_sessions_dir()?;
|
let db_path = Self::default_database_path()?;
|
||||||
Self::with_directory(sessions_dir)
|
Self::with_database_path(db_path).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create a storage manager with a custom sessions directory
|
/// Create a storage manager using the provided database path
|
||||||
pub fn with_directory(sessions_dir: PathBuf) -> Result<Self> {
|
pub async fn with_database_path(database_path: PathBuf) -> Result<Self> {
|
||||||
// Ensure the directory exists
|
if let Some(parent) = database_path.parent() {
|
||||||
if !sessions_dir.exists() {
|
if !parent.exists() {
|
||||||
fs::create_dir_all(&sessions_dir).map_err(|e| {
|
std::fs::create_dir_all(parent).map_err(|e| {
|
||||||
Error::Storage(format!("Failed to create sessions directory: {}", e))
|
Error::Storage(format!(
|
||||||
})?;
|
"Failed to create database directory {parent:?}: {e}"
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(Self { sessions_dir })
|
let options = SqliteConnectOptions::from_str(&format!(
|
||||||
|
"sqlite://{}",
|
||||||
|
database_path
|
||||||
|
.to_str()
|
||||||
|
.ok_or_else(|| Error::Storage("Invalid database path".to_string()))?
|
||||||
|
))
|
||||||
|
.map_err(|e| Error::Storage(format!("Invalid database URL: {e}")))?
|
||||||
|
.create_if_missing(true)
|
||||||
|
.journal_mode(SqliteJournalMode::Wal)
|
||||||
|
.synchronous(SqliteSynchronous::Normal);
|
||||||
|
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(5)
|
||||||
|
.connect_with(options)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to connect to database: {e}")))?;
|
||||||
|
|
||||||
|
sqlx::migrate!("./migrations")
|
||||||
|
.run(&pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to run database migrations: {e}")))?;
|
||||||
|
|
||||||
|
let storage = Self {
|
||||||
|
pool,
|
||||||
|
database_path,
|
||||||
|
};
|
||||||
|
|
||||||
|
storage.try_migrate_legacy_sessions().await?;
|
||||||
|
|
||||||
|
Ok(storage)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get the default sessions directory
|
/// Save a conversation. Existing entries are updated in-place.
|
||||||
/// - Linux: ~/.local/share/owlen/sessions
|
pub async fn save_conversation(
|
||||||
/// - Windows: %APPDATA%\owlen\sessions
|
&self,
|
||||||
/// - macOS: ~/Library/Application Support/owlen/sessions
|
conversation: &Conversation,
|
||||||
pub fn default_sessions_dir() -> Result<PathBuf> {
|
name: Option<String>,
|
||||||
|
) -> Result<()> {
|
||||||
|
self.save_conversation_with_description(conversation, name, None)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Save a conversation with an optional description override
|
||||||
|
pub async fn save_conversation_with_description(
|
||||||
|
&self,
|
||||||
|
conversation: &Conversation,
|
||||||
|
name: Option<String>,
|
||||||
|
description: Option<String>,
|
||||||
|
) -> Result<()> {
|
||||||
|
let mut serialized = conversation.clone();
|
||||||
|
if name.is_some() {
|
||||||
|
serialized.name = name.clone();
|
||||||
|
}
|
||||||
|
if description.is_some() {
|
||||||
|
serialized.description = description.clone();
|
||||||
|
}
|
||||||
|
|
||||||
|
let data = serde_json::to_string(&serialized)
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to serialize conversation: {e}")))?;
|
||||||
|
|
||||||
|
let created_at = to_epoch_seconds(serialized.created_at);
|
||||||
|
let updated_at = to_epoch_seconds(serialized.updated_at);
|
||||||
|
let message_count = serialized.messages.len() as i64;
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
r#"
|
||||||
|
INSERT INTO conversations (
|
||||||
|
id,
|
||||||
|
name,
|
||||||
|
description,
|
||||||
|
model,
|
||||||
|
message_count,
|
||||||
|
created_at,
|
||||||
|
updated_at,
|
||||||
|
data
|
||||||
|
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
|
||||||
|
ON CONFLICT(id) DO UPDATE SET
|
||||||
|
name = excluded.name,
|
||||||
|
description = excluded.description,
|
||||||
|
model = excluded.model,
|
||||||
|
message_count = excluded.message_count,
|
||||||
|
created_at = excluded.created_at,
|
||||||
|
updated_at = excluded.updated_at,
|
||||||
|
data = excluded.data
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(serialized.id.to_string())
|
||||||
|
.bind(name.or(serialized.name.clone()))
|
||||||
|
.bind(description.or(serialized.description.clone()))
|
||||||
|
.bind(&serialized.model)
|
||||||
|
.bind(message_count)
|
||||||
|
.bind(created_at)
|
||||||
|
.bind(updated_at)
|
||||||
|
.bind(data)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to save conversation: {e}")))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Load a conversation by ID
|
||||||
|
pub async fn load_conversation(&self, id: Uuid) -> Result<Conversation> {
|
||||||
|
let record = sqlx::query(r#"SELECT data FROM conversations WHERE id = ?1"#)
|
||||||
|
.bind(id.to_string())
|
||||||
|
.fetch_optional(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to load conversation: {e}")))?;
|
||||||
|
|
||||||
|
let row =
|
||||||
|
record.ok_or_else(|| Error::Storage(format!("No conversation found with id {id}")))?;
|
||||||
|
|
||||||
|
let data: String = row
|
||||||
|
.try_get("data")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read conversation payload: {e}")))?;
|
||||||
|
|
||||||
|
serde_json::from_str(&data)
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to deserialize conversation: {e}")))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List metadata for all saved conversations ordered by most recent update
|
||||||
|
pub async fn list_sessions(&self) -> Result<Vec<SessionMeta>> {
|
||||||
|
let rows = sqlx::query(
|
||||||
|
r#"
|
||||||
|
SELECT id, name, description, model, message_count, created_at, updated_at
|
||||||
|
FROM conversations
|
||||||
|
ORDER BY updated_at DESC
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to list sessions: {e}")))?;
|
||||||
|
|
||||||
|
let mut sessions = Vec::with_capacity(rows.len());
|
||||||
|
for row in rows {
|
||||||
|
let id_text: String = row
|
||||||
|
.try_get("id")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read id column: {e}")))?;
|
||||||
|
let id = Uuid::parse_str(&id_text)
|
||||||
|
.map_err(|e| Error::Storage(format!("Invalid UUID in storage: {e}")))?;
|
||||||
|
|
||||||
|
let message_count: i64 = row
|
||||||
|
.try_get("message_count")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read message count: {e}")))?;
|
||||||
|
|
||||||
|
let created_at: i64 = row
|
||||||
|
.try_get("created_at")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read created_at: {e}")))?;
|
||||||
|
let updated_at: i64 = row
|
||||||
|
.try_get("updated_at")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read updated_at: {e}")))?;
|
||||||
|
|
||||||
|
sessions.push(SessionMeta {
|
||||||
|
id,
|
||||||
|
name: row
|
||||||
|
.try_get("name")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read name: {e}")))?,
|
||||||
|
description: row
|
||||||
|
.try_get("description")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read description: {e}")))?,
|
||||||
|
model: row
|
||||||
|
.try_get("model")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read model: {e}")))?,
|
||||||
|
message_count: message_count as usize,
|
||||||
|
created_at: from_epoch_seconds(created_at),
|
||||||
|
updated_at: from_epoch_seconds(updated_at),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(sessions)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Delete a conversation by ID
|
||||||
|
pub async fn delete_session(&self, id: Uuid) -> Result<()> {
|
||||||
|
sqlx::query("DELETE FROM conversations WHERE id = ?1")
|
||||||
|
.bind(id.to_string())
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to delete conversation: {e}")))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn store_secure_item(
|
||||||
|
&self,
|
||||||
|
key: &str,
|
||||||
|
plaintext: &[u8],
|
||||||
|
master_key: &[u8],
|
||||||
|
) -> Result<()> {
|
||||||
|
let cipher = create_cipher(master_key)?;
|
||||||
|
let nonce_bytes = generate_nonce()?;
|
||||||
|
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||||
|
let ciphertext = cipher
|
||||||
|
.encrypt(nonce, plaintext)
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to encrypt secure item: {e}")))?;
|
||||||
|
|
||||||
|
let now = to_epoch_seconds(SystemTime::now());
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
r#"
|
||||||
|
INSERT INTO secure_items (key, nonce, ciphertext, created_at, updated_at)
|
||||||
|
VALUES (?1, ?2, ?3, ?4, ?5)
|
||||||
|
ON CONFLICT(key) DO UPDATE SET
|
||||||
|
nonce = excluded.nonce,
|
||||||
|
ciphertext = excluded.ciphertext,
|
||||||
|
updated_at = excluded.updated_at
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(key)
|
||||||
|
.bind(&nonce_bytes[..])
|
||||||
|
.bind(&ciphertext[..])
|
||||||
|
.bind(now)
|
||||||
|
.bind(now)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to store secure item: {e}")))?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn load_secure_item(&self, key: &str, master_key: &[u8]) -> Result<Option<Vec<u8>>> {
|
||||||
|
let record = sqlx::query("SELECT nonce, ciphertext FROM secure_items WHERE key = ?1")
|
||||||
|
.bind(key)
|
||||||
|
.fetch_optional(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to load secure item: {e}")))?;
|
||||||
|
|
||||||
|
let Some(row) = record else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
|
||||||
|
let nonce_bytes: Vec<u8> = row
|
||||||
|
.try_get("nonce")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read secure item nonce: {e}")))?;
|
||||||
|
let ciphertext: Vec<u8> = row
|
||||||
|
.try_get("ciphertext")
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read secure item ciphertext: {e}")))?;
|
||||||
|
|
||||||
|
if nonce_bytes.len() != 12 {
|
||||||
|
return Err(Error::Storage(
|
||||||
|
"Invalid nonce length for secure item".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let cipher = create_cipher(master_key)?;
|
||||||
|
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||||
|
let plaintext = cipher
|
||||||
|
.decrypt(nonce, ciphertext.as_ref())
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to decrypt secure item: {e}")))?;
|
||||||
|
|
||||||
|
Ok(Some(plaintext))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn delete_secure_item(&self, key: &str) -> Result<()> {
|
||||||
|
sqlx::query("DELETE FROM secure_items WHERE key = ?1")
|
||||||
|
.bind(key)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to delete secure item: {e}")))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn clear_secure_items(&self) -> Result<()> {
|
||||||
|
sqlx::query("DELETE FROM secure_items")
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to clear secure items: {e}")))?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Database location used by this storage manager
|
||||||
|
pub fn database_path(&self) -> &Path {
|
||||||
|
&self.database_path
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Determine default database path (platform specific)
|
||||||
|
pub fn default_database_path() -> Result<PathBuf> {
|
||||||
|
let data_dir = dirs::data_local_dir()
|
||||||
|
.ok_or_else(|| Error::Storage("Could not determine data directory".to_string()))?;
|
||||||
|
Ok(data_dir.join("owlen").join("owlen.db"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn legacy_sessions_dir() -> Result<PathBuf> {
|
||||||
let data_dir = dirs::data_local_dir()
|
let data_dir = dirs::data_local_dir()
|
||||||
.ok_or_else(|| Error::Storage("Could not determine data directory".to_string()))?;
|
.ok_or_else(|| Error::Storage("Could not determine data directory".to_string()))?;
|
||||||
Ok(data_dir.join("owlen").join("sessions"))
|
Ok(data_dir.join("owlen").join("sessions"))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Save a conversation to disk
|
async fn database_has_records(&self) -> Result<bool> {
|
||||||
pub fn save_conversation(
|
let (count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM conversations")
|
||||||
&self,
|
.fetch_one(&self.pool)
|
||||||
conversation: &Conversation,
|
.await
|
||||||
name: Option<String>,
|
.map_err(|e| Error::Storage(format!("Failed to inspect database: {e}")))?;
|
||||||
) -> Result<PathBuf> {
|
Ok(count > 0)
|
||||||
self.save_conversation_with_description(conversation, name, None)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Save a conversation to disk with an optional description
|
async fn try_migrate_legacy_sessions(&self) -> Result<()> {
|
||||||
pub fn save_conversation_with_description(
|
if self.database_has_records().await? {
|
||||||
&self,
|
return Ok(());
|
||||||
conversation: &Conversation,
|
}
|
||||||
name: Option<String>,
|
|
||||||
description: Option<String>,
|
let legacy_dir = match Self::legacy_sessions_dir() {
|
||||||
) -> Result<PathBuf> {
|
Ok(dir) => dir,
|
||||||
let filename = if let Some(ref session_name) = name {
|
Err(_) => return Ok(()),
|
||||||
// Use provided name, sanitized
|
|
||||||
let sanitized = sanitize_filename(session_name);
|
|
||||||
format!("{}_{}.json", conversation.id, sanitized)
|
|
||||||
} else {
|
|
||||||
// Use conversation ID and timestamp
|
|
||||||
let timestamp = SystemTime::now()
|
|
||||||
.duration_since(SystemTime::UNIX_EPOCH)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.as_secs();
|
|
||||||
format!("{}_{}.json", conversation.id, timestamp)
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let path = self.sessions_dir.join(filename);
|
if !legacy_dir.exists() {
|
||||||
|
return Ok(());
|
||||||
// Create a saveable version with the name and description
|
|
||||||
let mut save_conv = conversation.clone();
|
|
||||||
if name.is_some() {
|
|
||||||
save_conv.name = name;
|
|
||||||
}
|
|
||||||
if description.is_some() {
|
|
||||||
save_conv.description = description;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
let json = serde_json::to_string_pretty(&save_conv)
|
let entries = fs::read_dir(&legacy_dir).map_err(|e| {
|
||||||
.map_err(|e| Error::Storage(format!("Failed to serialize conversation: {}", e)))?;
|
Error::Storage(format!("Failed to read legacy sessions directory: {e}"))
|
||||||
|
})?;
|
||||||
fs::write(&path, json)
|
|
||||||
.map_err(|e| Error::Storage(format!("Failed to write session file: {}", e)))?;
|
|
||||||
|
|
||||||
Ok(path)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Load a conversation from disk
|
|
||||||
pub fn load_conversation(&self, path: impl AsRef<Path>) -> Result<Conversation> {
|
|
||||||
let content = fs::read_to_string(path.as_ref())
|
|
||||||
.map_err(|e| Error::Storage(format!("Failed to read session file: {}", e)))?;
|
|
||||||
|
|
||||||
let conversation: Conversation = serde_json::from_str(&content)
|
|
||||||
.map_err(|e| Error::Storage(format!("Failed to parse session file: {}", e)))?;
|
|
||||||
|
|
||||||
Ok(conversation)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// List all saved sessions with metadata
|
|
||||||
pub fn list_sessions(&self) -> Result<Vec<SessionMeta>> {
|
|
||||||
let mut sessions = Vec::new();
|
|
||||||
|
|
||||||
let entries = fs::read_dir(&self.sessions_dir)
|
|
||||||
.map_err(|e| Error::Storage(format!("Failed to read sessions directory: {}", e)))?;
|
|
||||||
|
|
||||||
for entry in entries {
|
|
||||||
let entry = entry
|
|
||||||
.map_err(|e| Error::Storage(format!("Failed to read directory entry: {}", e)))?;
|
|
||||||
|
|
||||||
|
let mut json_files = Vec::new();
|
||||||
|
for entry in entries.flatten() {
|
||||||
let path = entry.path();
|
let path = entry.path();
|
||||||
if path.extension().and_then(|s| s.to_str()) != Some("json") {
|
if path.extension().and_then(|s| s.to_str()) == Some("json") {
|
||||||
continue;
|
json_files.push(path);
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Try to load the conversation to extract metadata
|
if json_files.is_empty() {
|
||||||
match self.load_conversation(&path) {
|
return Ok(());
|
||||||
Ok(conv) => {
|
}
|
||||||
sessions.push(SessionMeta {
|
|
||||||
path: path.clone(),
|
if !io::stdin().is_terminal() {
|
||||||
id: conv.id,
|
return Ok(());
|
||||||
name: conv.name.clone(),
|
}
|
||||||
description: conv.description.clone(),
|
|
||||||
message_count: conv.messages.len(),
|
println!(
|
||||||
model: conv.model.clone(),
|
"Legacy OWLEN session files were found in {}.",
|
||||||
created_at: conv.created_at,
|
legacy_dir.display()
|
||||||
updated_at: conv.updated_at,
|
);
|
||||||
});
|
if !prompt_yes_no("Migrate them to the new SQLite storage? (y/N) ")? {
|
||||||
}
|
println!("Skipping legacy session migration.");
|
||||||
Err(_) => {
|
return Ok(());
|
||||||
// Skip files that can't be parsed
|
}
|
||||||
continue;
|
|
||||||
|
println!("Migrating legacy sessions...");
|
||||||
|
let mut migrated = 0usize;
|
||||||
|
for path in &json_files {
|
||||||
|
match fs::read_to_string(path) {
|
||||||
|
Ok(content) => match serde_json::from_str::<Conversation>(&content) {
|
||||||
|
Ok(conversation) => {
|
||||||
|
if let Err(err) = self
|
||||||
|
.save_conversation_with_description(
|
||||||
|
&conversation,
|
||||||
|
conversation.name.clone(),
|
||||||
|
conversation.description.clone(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
println!(" • Failed to migrate {}: {}", path.display(), err);
|
||||||
|
} else {
|
||||||
|
migrated += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
println!(
|
||||||
|
" • Failed to parse conversation {}: {}",
|
||||||
|
path.display(),
|
||||||
|
err
|
||||||
|
);
|
||||||
|
}
|
||||||
|
},
|
||||||
|
Err(err) => {
|
||||||
|
println!(" • Failed to read {}: {}", path.display(), err);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Sort by updated_at, most recent first
|
if migrated > 0 {
|
||||||
sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
|
if let Err(err) = archive_legacy_directory(&legacy_dir) {
|
||||||
|
println!(
|
||||||
Ok(sessions)
|
"Warning: migrated sessions but failed to archive legacy directory: {}",
|
||||||
}
|
err
|
||||||
|
);
|
||||||
/// Delete a saved session
|
|
||||||
pub fn delete_session(&self, path: impl AsRef<Path>) -> Result<()> {
|
|
||||||
fs::remove_file(path.as_ref())
|
|
||||||
.map_err(|e| Error::Storage(format!("Failed to delete session file: {}", e)))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Get the sessions directory path
|
|
||||||
pub fn sessions_dir(&self) -> &Path {
|
|
||||||
&self.sessions_dir
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Default for StorageManager {
|
|
||||||
fn default() -> Self {
|
|
||||||
Self::new().expect("Failed to create default storage manager")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Sanitize a filename by removing invalid characters
|
|
||||||
fn sanitize_filename(name: &str) -> String {
|
|
||||||
name.chars()
|
|
||||||
.map(|c| {
|
|
||||||
if c.is_alphanumeric() || c == '_' || c == '-' {
|
|
||||||
c
|
|
||||||
} else if c.is_whitespace() {
|
|
||||||
'_'
|
|
||||||
} else {
|
|
||||||
'-'
|
|
||||||
}
|
}
|
||||||
})
|
}
|
||||||
.collect::<String>()
|
|
||||||
.chars()
|
println!("Migrated {} legacy sessions.", migrated);
|
||||||
.take(50) // Limit length
|
Ok(())
|
||||||
.collect()
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn to_epoch_seconds(time: SystemTime) -> i64 {
|
||||||
|
match time.duration_since(UNIX_EPOCH) {
|
||||||
|
Ok(duration) => duration.as_secs() as i64,
|
||||||
|
Err(_) => 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn from_epoch_seconds(seconds: i64) -> SystemTime {
|
||||||
|
UNIX_EPOCH + Duration::from_secs(seconds.max(0) as u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn prompt_yes_no(prompt: &str) -> Result<bool> {
|
||||||
|
print!("{}", prompt);
|
||||||
|
io::stdout()
|
||||||
|
.flush()
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to flush stdout: {e}")))?;
|
||||||
|
|
||||||
|
let mut input = String::new();
|
||||||
|
io::stdin()
|
||||||
|
.read_line(&mut input)
|
||||||
|
.map_err(|e| Error::Storage(format!("Failed to read input: {e}")))?;
|
||||||
|
let trimmed = input.trim().to_lowercase();
|
||||||
|
Ok(matches!(trimmed.as_str(), "y" | "yes"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn archive_legacy_directory(legacy_dir: &Path) -> Result<()> {
|
||||||
|
let mut backup_dir = legacy_dir.with_file_name("sessions_legacy_backup");
|
||||||
|
let mut counter = 1;
|
||||||
|
while backup_dir.exists() {
|
||||||
|
backup_dir = legacy_dir.with_file_name(format!("sessions_legacy_backup_{}", counter));
|
||||||
|
counter += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
fs::rename(legacy_dir, &backup_dir).map_err(|e| {
|
||||||
|
Error::Storage(format!(
|
||||||
|
"Failed to archive legacy sessions directory {}: {}",
|
||||||
|
legacy_dir.display(),
|
||||||
|
e
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
println!("Legacy session files archived to {}", backup_dir.display());
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn create_cipher(master_key: &[u8]) -> Result<Aes256Gcm> {
|
||||||
|
if master_key.len() != 32 {
|
||||||
|
return Err(Error::Storage(
|
||||||
|
"Master key must be 32 bytes for AES-256-GCM".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Aes256Gcm::new_from_slice(master_key).map_err(|_| {
|
||||||
|
Error::Storage("Failed to initialize cipher with provided master key".to_string())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn generate_nonce() -> Result<[u8; 12]> {
|
||||||
|
let mut nonce = [0u8; 12];
|
||||||
|
SystemRandom::new()
|
||||||
|
.fill(&mut nonce)
|
||||||
|
.map_err(|_| Error::Storage("Failed to generate nonce".to_string()))?;
|
||||||
|
Ok(nonce)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::types::Message;
|
use crate::types::{Conversation, Message};
|
||||||
use tempfile::TempDir;
|
use tempfile::tempdir;
|
||||||
|
|
||||||
#[test]
|
fn sample_conversation() -> Conversation {
|
||||||
fn test_platform_specific_default_path() {
|
Conversation {
|
||||||
let path = StorageManager::default_sessions_dir().unwrap();
|
id: Uuid::new_v4(),
|
||||||
|
name: Some("Test conversation".to_string()),
|
||||||
// Verify it contains owlen/sessions
|
description: Some("A sample conversation".to_string()),
|
||||||
assert!(path.to_string_lossy().contains("owlen"));
|
messages: vec![
|
||||||
assert!(path.to_string_lossy().contains("sessions"));
|
Message::user("Hello".to_string()),
|
||||||
|
Message::assistant("Hi".to_string()),
|
||||||
// Platform-specific checks
|
],
|
||||||
#[cfg(target_os = "linux")]
|
model: "test-model".to_string(),
|
||||||
{
|
created_at: SystemTime::now(),
|
||||||
// Linux should use ~/.local/share/owlen/sessions
|
updated_at: SystemTime::now(),
|
||||||
assert!(path.to_string_lossy().contains(".local/share"));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(target_os = "windows")]
|
|
||||||
{
|
|
||||||
// Windows should use AppData
|
|
||||||
assert!(path.to_string_lossy().contains("AppData"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(target_os = "macos")]
|
|
||||||
{
|
|
||||||
// macOS should use ~/Library/Application Support
|
|
||||||
assert!(path
|
|
||||||
.to_string_lossy()
|
|
||||||
.contains("Library/Application Support"));
|
|
||||||
}
|
|
||||||
|
|
||||||
println!("Default sessions directory: {}", path.display());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
fn test_sanitize_filename() {
|
async fn test_storage_lifecycle() {
|
||||||
assert_eq!(sanitize_filename("Hello World"), "Hello_World");
|
let temp_dir = tempdir().expect("failed to create temp dir");
|
||||||
assert_eq!(sanitize_filename("test/path\\file"), "test-path-file");
|
let db_path = temp_dir.path().join("owlen.db");
|
||||||
assert_eq!(sanitize_filename("file:name?"), "file-name-");
|
let storage = StorageManager::with_database_path(db_path).await.unwrap();
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
let conversation = sample_conversation();
|
||||||
fn test_save_and_load_conversation() {
|
storage
|
||||||
let temp_dir = TempDir::new().unwrap();
|
.save_conversation(&conversation, None)
|
||||||
let storage = StorageManager::with_directory(temp_dir.path().to_path_buf()).unwrap();
|
.await
|
||||||
|
.expect("failed to save conversation");
|
||||||
|
|
||||||
let mut conv = Conversation::new("test-model".to_string());
|
let sessions = storage.list_sessions().await.unwrap();
|
||||||
conv.messages.push(Message::user("Hello".to_string()));
|
assert_eq!(sessions.len(), 1);
|
||||||
conv.messages
|
assert_eq!(sessions[0].id, conversation.id);
|
||||||
.push(Message::assistant("Hi there!".to_string()));
|
|
||||||
|
|
||||||
// Save conversation
|
let loaded = storage.load_conversation(conversation.id).await.unwrap();
|
||||||
let path = storage
|
|
||||||
.save_conversation(&conv, Some("test_session".to_string()))
|
|
||||||
.unwrap();
|
|
||||||
assert!(path.exists());
|
|
||||||
|
|
||||||
// Load conversation
|
|
||||||
let loaded = storage.load_conversation(&path).unwrap();
|
|
||||||
assert_eq!(loaded.id, conv.id);
|
|
||||||
assert_eq!(loaded.model, conv.model);
|
|
||||||
assert_eq!(loaded.messages.len(), 2);
|
assert_eq!(loaded.messages.len(), 2);
|
||||||
assert_eq!(loaded.name, Some("test_session".to_string()));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
storage
|
||||||
fn test_list_sessions() {
|
.delete_session(conversation.id)
|
||||||
let temp_dir = TempDir::new().unwrap();
|
.await
|
||||||
let storage = StorageManager::with_directory(temp_dir.path().to_path_buf()).unwrap();
|
.expect("failed to delete conversation");
|
||||||
|
let sessions = storage.list_sessions().await.unwrap();
|
||||||
// Create multiple sessions
|
assert!(sessions.is_empty());
|
||||||
for i in 0..3 {
|
|
||||||
let mut conv = Conversation::new("test-model".to_string());
|
|
||||||
conv.messages.push(Message::user(format!("Message {}", i)));
|
|
||||||
storage
|
|
||||||
.save_conversation(&conv, Some(format!("session_{}", i)))
|
|
||||||
.unwrap();
|
|
||||||
}
|
|
||||||
|
|
||||||
// List sessions
|
|
||||||
let sessions = storage.list_sessions().unwrap();
|
|
||||||
assert_eq!(sessions.len(), 3);
|
|
||||||
|
|
||||||
// Check that sessions are sorted by updated_at (most recent first)
|
|
||||||
for i in 0..sessions.len() - 1 {
|
|
||||||
assert!(sessions[i].updated_at >= sessions[i + 1].updated_at);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_delete_session() {
|
|
||||||
let temp_dir = TempDir::new().unwrap();
|
|
||||||
let storage = StorageManager::with_directory(temp_dir.path().to_path_buf()).unwrap();
|
|
||||||
|
|
||||||
let conv = Conversation::new("test-model".to_string());
|
|
||||||
let path = storage.save_conversation(&conv, None).unwrap();
|
|
||||||
assert!(path.exists());
|
|
||||||
|
|
||||||
storage.delete_session(&path).unwrap();
|
|
||||||
assert!(!path.exists());
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -44,6 +44,11 @@ pub struct Theme {
|
|||||||
#[serde(serialize_with = "serialize_color")]
|
#[serde(serialize_with = "serialize_color")]
|
||||||
pub assistant_message_role: Color,
|
pub assistant_message_role: Color,
|
||||||
|
|
||||||
|
/// Color for tool output messages
|
||||||
|
#[serde(deserialize_with = "deserialize_color")]
|
||||||
|
#[serde(serialize_with = "serialize_color")]
|
||||||
|
pub tool_output: Color,
|
||||||
|
|
||||||
/// Color for thinking panel title
|
/// Color for thinking panel title
|
||||||
#[serde(deserialize_with = "deserialize_color")]
|
#[serde(deserialize_with = "deserialize_color")]
|
||||||
#[serde(serialize_with = "serialize_color")]
|
#[serde(serialize_with = "serialize_color")]
|
||||||
@@ -268,6 +273,7 @@ fn default_dark() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(95, 20, 135),
|
unfocused_panel_border: Color::Rgb(95, 20, 135),
|
||||||
user_message_role: Color::LightBlue,
|
user_message_role: Color::LightBlue,
|
||||||
assistant_message_role: Color::Yellow,
|
assistant_message_role: Color::Yellow,
|
||||||
|
tool_output: Color::Gray,
|
||||||
thinking_panel_title: Color::LightMagenta,
|
thinking_panel_title: Color::LightMagenta,
|
||||||
command_bar_background: Color::Black,
|
command_bar_background: Color::Black,
|
||||||
status_background: Color::Black,
|
status_background: Color::Black,
|
||||||
@@ -297,6 +303,7 @@ fn default_light() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(221, 221, 221),
|
unfocused_panel_border: Color::Rgb(221, 221, 221),
|
||||||
user_message_role: Color::Rgb(0, 85, 164),
|
user_message_role: Color::Rgb(0, 85, 164),
|
||||||
assistant_message_role: Color::Rgb(142, 68, 173),
|
assistant_message_role: Color::Rgb(142, 68, 173),
|
||||||
|
tool_output: Color::Gray,
|
||||||
thinking_panel_title: Color::Rgb(142, 68, 173),
|
thinking_panel_title: Color::Rgb(142, 68, 173),
|
||||||
command_bar_background: Color::White,
|
command_bar_background: Color::White,
|
||||||
status_background: Color::White,
|
status_background: Color::White,
|
||||||
@@ -326,8 +333,9 @@ fn gruvbox() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(124, 111, 100), // #7c6f64
|
unfocused_panel_border: Color::Rgb(124, 111, 100), // #7c6f64
|
||||||
user_message_role: Color::Rgb(184, 187, 38), // #b8bb26 (green)
|
user_message_role: Color::Rgb(184, 187, 38), // #b8bb26 (green)
|
||||||
assistant_message_role: Color::Rgb(131, 165, 152), // #83a598 (blue)
|
assistant_message_role: Color::Rgb(131, 165, 152), // #83a598 (blue)
|
||||||
thinking_panel_title: Color::Rgb(211, 134, 155), // #d3869b (purple)
|
tool_output: Color::Rgb(146, 131, 116),
|
||||||
command_bar_background: Color::Rgb(60, 56, 54), // #3c3836
|
thinking_panel_title: Color::Rgb(211, 134, 155), // #d3869b (purple)
|
||||||
|
command_bar_background: Color::Rgb(60, 56, 54), // #3c3836
|
||||||
status_background: Color::Rgb(60, 56, 54),
|
status_background: Color::Rgb(60, 56, 54),
|
||||||
mode_normal: Color::Rgb(131, 165, 152), // blue
|
mode_normal: Color::Rgb(131, 165, 152), // blue
|
||||||
mode_editing: Color::Rgb(184, 187, 38), // green
|
mode_editing: Color::Rgb(184, 187, 38), // green
|
||||||
@@ -355,7 +363,8 @@ fn dracula() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(68, 71, 90), // #44475a
|
unfocused_panel_border: Color::Rgb(68, 71, 90), // #44475a
|
||||||
user_message_role: Color::Rgb(139, 233, 253), // #8be9fd (cyan)
|
user_message_role: Color::Rgb(139, 233, 253), // #8be9fd (cyan)
|
||||||
assistant_message_role: Color::Rgb(255, 121, 198), // #ff79c6 (pink)
|
assistant_message_role: Color::Rgb(255, 121, 198), // #ff79c6 (pink)
|
||||||
thinking_panel_title: Color::Rgb(189, 147, 249), // #bd93f9 (purple)
|
tool_output: Color::Rgb(98, 114, 164),
|
||||||
|
thinking_panel_title: Color::Rgb(189, 147, 249), // #bd93f9 (purple)
|
||||||
command_bar_background: Color::Rgb(68, 71, 90),
|
command_bar_background: Color::Rgb(68, 71, 90),
|
||||||
status_background: Color::Rgb(68, 71, 90),
|
status_background: Color::Rgb(68, 71, 90),
|
||||||
mode_normal: Color::Rgb(139, 233, 253),
|
mode_normal: Color::Rgb(139, 233, 253),
|
||||||
@@ -384,6 +393,7 @@ fn solarized() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(7, 54, 66), // #073642 (base02)
|
unfocused_panel_border: Color::Rgb(7, 54, 66), // #073642 (base02)
|
||||||
user_message_role: Color::Rgb(42, 161, 152), // #2aa198 (cyan)
|
user_message_role: Color::Rgb(42, 161, 152), // #2aa198 (cyan)
|
||||||
assistant_message_role: Color::Rgb(203, 75, 22), // #cb4b16 (orange)
|
assistant_message_role: Color::Rgb(203, 75, 22), // #cb4b16 (orange)
|
||||||
|
tool_output: Color::Rgb(101, 123, 131),
|
||||||
thinking_panel_title: Color::Rgb(108, 113, 196), // #6c71c4 (violet)
|
thinking_panel_title: Color::Rgb(108, 113, 196), // #6c71c4 (violet)
|
||||||
command_bar_background: Color::Rgb(7, 54, 66),
|
command_bar_background: Color::Rgb(7, 54, 66),
|
||||||
status_background: Color::Rgb(7, 54, 66),
|
status_background: Color::Rgb(7, 54, 66),
|
||||||
@@ -413,6 +423,7 @@ fn midnight_ocean() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(48, 54, 61),
|
unfocused_panel_border: Color::Rgb(48, 54, 61),
|
||||||
user_message_role: Color::Rgb(121, 192, 255),
|
user_message_role: Color::Rgb(121, 192, 255),
|
||||||
assistant_message_role: Color::Rgb(137, 221, 255),
|
assistant_message_role: Color::Rgb(137, 221, 255),
|
||||||
|
tool_output: Color::Rgb(84, 110, 122),
|
||||||
thinking_panel_title: Color::Rgb(158, 206, 106),
|
thinking_panel_title: Color::Rgb(158, 206, 106),
|
||||||
command_bar_background: Color::Rgb(22, 27, 34),
|
command_bar_background: Color::Rgb(22, 27, 34),
|
||||||
status_background: Color::Rgb(22, 27, 34),
|
status_background: Color::Rgb(22, 27, 34),
|
||||||
@@ -442,7 +453,8 @@ fn rose_pine() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(38, 35, 58), // #26233a
|
unfocused_panel_border: Color::Rgb(38, 35, 58), // #26233a
|
||||||
user_message_role: Color::Rgb(49, 116, 143), // #31748f (foam)
|
user_message_role: Color::Rgb(49, 116, 143), // #31748f (foam)
|
||||||
assistant_message_role: Color::Rgb(156, 207, 216), // #9ccfd8 (foam light)
|
assistant_message_role: Color::Rgb(156, 207, 216), // #9ccfd8 (foam light)
|
||||||
thinking_panel_title: Color::Rgb(196, 167, 231), // #c4a7e7 (iris)
|
tool_output: Color::Rgb(110, 106, 134),
|
||||||
|
thinking_panel_title: Color::Rgb(196, 167, 231), // #c4a7e7 (iris)
|
||||||
command_bar_background: Color::Rgb(38, 35, 58),
|
command_bar_background: Color::Rgb(38, 35, 58),
|
||||||
status_background: Color::Rgb(38, 35, 58),
|
status_background: Color::Rgb(38, 35, 58),
|
||||||
mode_normal: Color::Rgb(156, 207, 216),
|
mode_normal: Color::Rgb(156, 207, 216),
|
||||||
@@ -471,7 +483,8 @@ fn monokai() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(117, 113, 94), // #75715e
|
unfocused_panel_border: Color::Rgb(117, 113, 94), // #75715e
|
||||||
user_message_role: Color::Rgb(102, 217, 239), // #66d9ef (cyan)
|
user_message_role: Color::Rgb(102, 217, 239), // #66d9ef (cyan)
|
||||||
assistant_message_role: Color::Rgb(174, 129, 255), // #ae81ff (purple)
|
assistant_message_role: Color::Rgb(174, 129, 255), // #ae81ff (purple)
|
||||||
thinking_panel_title: Color::Rgb(230, 219, 116), // #e6db74 (yellow)
|
tool_output: Color::Rgb(117, 113, 94),
|
||||||
|
thinking_panel_title: Color::Rgb(230, 219, 116), // #e6db74 (yellow)
|
||||||
command_bar_background: Color::Rgb(39, 40, 34),
|
command_bar_background: Color::Rgb(39, 40, 34),
|
||||||
status_background: Color::Rgb(39, 40, 34),
|
status_background: Color::Rgb(39, 40, 34),
|
||||||
mode_normal: Color::Rgb(102, 217, 239),
|
mode_normal: Color::Rgb(102, 217, 239),
|
||||||
@@ -500,7 +513,8 @@ fn material_dark() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(84, 110, 122), // #546e7a
|
unfocused_panel_border: Color::Rgb(84, 110, 122), // #546e7a
|
||||||
user_message_role: Color::Rgb(130, 170, 255), // #82aaff (blue)
|
user_message_role: Color::Rgb(130, 170, 255), // #82aaff (blue)
|
||||||
assistant_message_role: Color::Rgb(199, 146, 234), // #c792ea (purple)
|
assistant_message_role: Color::Rgb(199, 146, 234), // #c792ea (purple)
|
||||||
thinking_panel_title: Color::Rgb(255, 203, 107), // #ffcb6b (yellow)
|
tool_output: Color::Rgb(84, 110, 122),
|
||||||
|
thinking_panel_title: Color::Rgb(255, 203, 107), // #ffcb6b (yellow)
|
||||||
command_bar_background: Color::Rgb(33, 43, 48),
|
command_bar_background: Color::Rgb(33, 43, 48),
|
||||||
status_background: Color::Rgb(33, 43, 48),
|
status_background: Color::Rgb(33, 43, 48),
|
||||||
mode_normal: Color::Rgb(130, 170, 255),
|
mode_normal: Color::Rgb(130, 170, 255),
|
||||||
@@ -529,6 +543,7 @@ fn material_light() -> Theme {
|
|||||||
unfocused_panel_border: Color::Rgb(176, 190, 197),
|
unfocused_panel_border: Color::Rgb(176, 190, 197),
|
||||||
user_message_role: Color::Rgb(68, 138, 255),
|
user_message_role: Color::Rgb(68, 138, 255),
|
||||||
assistant_message_role: Color::Rgb(124, 77, 255),
|
assistant_message_role: Color::Rgb(124, 77, 255),
|
||||||
|
tool_output: Color::Rgb(144, 164, 174),
|
||||||
thinking_panel_title: Color::Rgb(245, 124, 0),
|
thinking_panel_title: Color::Rgb(245, 124, 0),
|
||||||
command_bar_background: Color::Rgb(255, 255, 255),
|
command_bar_background: Color::Rgb(255, 255, 255),
|
||||||
status_background: Color::Rgb(255, 255, 255),
|
status_background: Color::Rgb(255, 255, 255),
|
||||||
|
|||||||
95
crates/owlen-core/src/tools.rs
Normal file
95
crates/owlen-core/src/tools.rs
Normal file
@@ -0,0 +1,95 @@
|
|||||||
|
//! Tool module aggregating built‑in tool implementations.
|
||||||
|
//!
|
||||||
|
//! The crate originally declared `pub mod tools;` in `lib.rs` but the source
|
||||||
|
//! directory only contained individual tool files without a `mod.rs`, causing the
|
||||||
|
//! compiler to look for `tools.rs` and fail. Adding this module file makes the
|
||||||
|
//! directory a proper Rust module and re‑exports the concrete tool types.
|
||||||
|
|
||||||
|
pub mod code_exec;
|
||||||
|
pub mod fs_tools;
|
||||||
|
pub mod registry;
|
||||||
|
pub mod web_search;
|
||||||
|
pub mod web_search_detailed;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use crate::Result;
|
||||||
|
|
||||||
|
/// Trait representing a tool that can be called via the MCP interface.
|
||||||
|
#[async_trait]
|
||||||
|
pub trait Tool: Send + Sync {
|
||||||
|
/// Unique name of the tool (used in the MCP protocol).
|
||||||
|
fn name(&self) -> &'static str;
|
||||||
|
/// Human‑readable description for documentation.
|
||||||
|
fn description(&self) -> &'static str;
|
||||||
|
/// JSON‑Schema describing the expected arguments.
|
||||||
|
fn schema(&self) -> Value;
|
||||||
|
/// Execute the tool with the provided arguments.
|
||||||
|
fn requires_network(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
fn requires_filesystem(&self) -> Vec<String> {
|
||||||
|
Vec::new()
|
||||||
|
}
|
||||||
|
async fn execute(&self, args: Value) -> Result<ToolResult>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result returned by a tool execution.
|
||||||
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct ToolResult {
|
||||||
|
/// Indicates whether the tool completed successfully.
|
||||||
|
pub success: bool,
|
||||||
|
/// Human‑readable status string – retained for compatibility.
|
||||||
|
pub status: String,
|
||||||
|
/// Arbitrary JSON payload describing the tool output.
|
||||||
|
pub output: Value,
|
||||||
|
/// Execution duration.
|
||||||
|
#[serde(skip_serializing_if = "Duration::is_zero", default)]
|
||||||
|
pub duration: Duration,
|
||||||
|
/// Optional key/value metadata for the tool invocation.
|
||||||
|
#[serde(default)]
|
||||||
|
pub metadata: HashMap<String, String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolResult {
|
||||||
|
pub fn success(output: Value) -> Self {
|
||||||
|
Self {
|
||||||
|
success: true,
|
||||||
|
status: "success".into(),
|
||||||
|
output,
|
||||||
|
duration: Duration::default(),
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn error(msg: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
success: false,
|
||||||
|
status: "error".into(),
|
||||||
|
output: json!({ "error": msg }),
|
||||||
|
duration: Duration::default(),
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn cancelled(msg: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
success: false,
|
||||||
|
status: "cancelled".into(),
|
||||||
|
output: json!({ "error": msg }),
|
||||||
|
duration: Duration::default(),
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Re‑export the most commonly used types so they can be accessed as
|
||||||
|
// `owlen_core::tools::CodeExecTool`, etc.
|
||||||
|
pub use code_exec::CodeExecTool;
|
||||||
|
pub use fs_tools::{ResourcesDeleteTool, ResourcesGetTool, ResourcesListTool, ResourcesWriteTool};
|
||||||
|
pub use registry::ToolRegistry;
|
||||||
|
pub use web_search::WebSearchTool;
|
||||||
|
pub use web_search_detailed::WebSearchDetailedTool;
|
||||||
148
crates/owlen-core/src/tools/code_exec.rs
Normal file
148
crates/owlen-core/src/tools/code_exec.rs
Normal file
@@ -0,0 +1,148 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use crate::Result;
|
||||||
|
use anyhow::{anyhow, Context};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
use super::{Tool, ToolResult};
|
||||||
|
use crate::sandbox::{SandboxConfig, SandboxedProcess};
|
||||||
|
|
||||||
|
pub struct CodeExecTool {
|
||||||
|
allowed_languages: Arc<Vec<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CodeExecTool {
|
||||||
|
pub fn new(allowed_languages: Vec<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
allowed_languages: Arc::new(allowed_languages),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for CodeExecTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"code_exec"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Execute code snippets within a sandboxed environment"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"language": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": self.allowed_languages.as_slice(),
|
||||||
|
"description": "Language of the code block"
|
||||||
|
},
|
||||||
|
"code": {
|
||||||
|
"type": "string",
|
||||||
|
"minLength": 1,
|
||||||
|
"maxLength": 10000,
|
||||||
|
"description": "Code to execute"
|
||||||
|
},
|
||||||
|
"timeout": {
|
||||||
|
"type": "integer",
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 300,
|
||||||
|
"default": 30,
|
||||||
|
"description": "Execution timeout in seconds"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["language", "code"],
|
||||||
|
"additionalProperties": false
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value) -> Result<ToolResult> {
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
let language = args
|
||||||
|
.get("language")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.context("Missing language parameter")?;
|
||||||
|
let code = args
|
||||||
|
.get("code")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.context("Missing code parameter")?;
|
||||||
|
let timeout = args.get("timeout").and_then(Value::as_u64).unwrap_or(30);
|
||||||
|
|
||||||
|
if !self.allowed_languages.iter().any(|lang| lang == language) {
|
||||||
|
return Err(anyhow!("Language '{}' not permitted", language).into());
|
||||||
|
}
|
||||||
|
|
||||||
|
let (command, command_args) = match language {
|
||||||
|
"python" => (
|
||||||
|
"python3".to_string(),
|
||||||
|
vec!["-c".to_string(), code.to_string()],
|
||||||
|
),
|
||||||
|
"javascript" => ("node".to_string(), vec!["-e".to_string(), code.to_string()]),
|
||||||
|
"bash" => ("bash".to_string(), vec!["-c".to_string(), code.to_string()]),
|
||||||
|
"rust" => {
|
||||||
|
let mut result =
|
||||||
|
ToolResult::error("Rust execution is not yet supported in the sandbox");
|
||||||
|
result.duration = start.elapsed();
|
||||||
|
return Ok(result);
|
||||||
|
}
|
||||||
|
other => return Err(anyhow!("Unsupported language: {}", other).into()),
|
||||||
|
};
|
||||||
|
|
||||||
|
let sandbox_config = SandboxConfig {
|
||||||
|
allow_network: false,
|
||||||
|
timeout_seconds: timeout,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let sandbox_result = tokio::task::spawn_blocking(move || {
|
||||||
|
let sandbox = SandboxedProcess::new(sandbox_config)?;
|
||||||
|
let arg_refs: Vec<&str> = command_args.iter().map(|s| s.as_str()).collect();
|
||||||
|
sandbox.execute(&command, &arg_refs)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.context("Sandbox execution task failed")??;
|
||||||
|
|
||||||
|
let mut result = if sandbox_result.exit_code == 0 {
|
||||||
|
ToolResult::success(json!({
|
||||||
|
"stdout": sandbox_result.stdout,
|
||||||
|
"stderr": sandbox_result.stderr,
|
||||||
|
"exit_code": sandbox_result.exit_code,
|
||||||
|
"timed_out": sandbox_result.was_timeout,
|
||||||
|
}))
|
||||||
|
} else {
|
||||||
|
let error_msg = if sandbox_result.was_timeout {
|
||||||
|
format!(
|
||||||
|
"Execution timed out after {} seconds (exit code {}): {}",
|
||||||
|
timeout, sandbox_result.exit_code, sandbox_result.stderr
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"Execution failed with status {}: {}",
|
||||||
|
sandbox_result.exit_code, sandbox_result.stderr
|
||||||
|
)
|
||||||
|
};
|
||||||
|
let mut err_result = ToolResult::error(&error_msg);
|
||||||
|
err_result.output = json!({
|
||||||
|
"stdout": sandbox_result.stdout,
|
||||||
|
"stderr": sandbox_result.stderr,
|
||||||
|
"exit_code": sandbox_result.exit_code,
|
||||||
|
"timed_out": sandbox_result.was_timeout,
|
||||||
|
});
|
||||||
|
err_result
|
||||||
|
};
|
||||||
|
|
||||||
|
result.duration = start.elapsed();
|
||||||
|
result
|
||||||
|
.metadata
|
||||||
|
.insert("language".to_string(), language.to_string());
|
||||||
|
result
|
||||||
|
.metadata
|
||||||
|
.insert("timeout_seconds".to_string(), timeout.to_string());
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
}
|
||||||
198
crates/owlen-core/src/tools/fs_tools.rs
Normal file
198
crates/owlen-core/src/tools/fs_tools.rs
Normal file
@@ -0,0 +1,198 @@
|
|||||||
|
use crate::tools::{Tool, ToolResult};
|
||||||
|
use crate::{Error, Result};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use path_clean::PathClean;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::env;
|
||||||
|
use std::fs;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct FileArgs {
|
||||||
|
path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sanitize_path(path: &str, root: &Path) -> Result<PathBuf> {
|
||||||
|
let path = Path::new(path);
|
||||||
|
let path = if path.is_absolute() {
|
||||||
|
// Strip leading '/' to treat as relative to the project root.
|
||||||
|
path.strip_prefix("/")
|
||||||
|
.map_err(|_| Error::InvalidInput("Invalid path".into()))?
|
||||||
|
.to_path_buf()
|
||||||
|
} else {
|
||||||
|
path.to_path_buf()
|
||||||
|
};
|
||||||
|
|
||||||
|
let full_path = root.join(path).clean();
|
||||||
|
|
||||||
|
if !full_path.starts_with(root) {
|
||||||
|
return Err(Error::PermissionDenied("Path traversal detected".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(full_path)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ResourcesListTool;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for ResourcesListTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"resources/list"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Lists directory contents."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The path to the directory to list."
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["path"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: serde_json::Value) -> Result<ToolResult> {
|
||||||
|
let args: FileArgs = serde_json::from_value(args)?;
|
||||||
|
let root = env::current_dir()?;
|
||||||
|
let full_path = sanitize_path(&args.path, &root)?;
|
||||||
|
|
||||||
|
let entries = fs::read_dir(full_path)?;
|
||||||
|
|
||||||
|
let mut result = Vec::new();
|
||||||
|
for entry in entries {
|
||||||
|
let entry = entry?;
|
||||||
|
result.push(entry.file_name().to_string_lossy().to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(ToolResult::success(serde_json::to_value(result)?))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ResourcesGetTool;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for ResourcesGetTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"resources/get"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Reads file content."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The path to the file to read."
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["path"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: serde_json::Value) -> Result<ToolResult> {
|
||||||
|
let args: FileArgs = serde_json::from_value(args)?;
|
||||||
|
let root = env::current_dir()?;
|
||||||
|
let full_path = sanitize_path(&args.path, &root)?;
|
||||||
|
|
||||||
|
let content = fs::read_to_string(full_path)?;
|
||||||
|
|
||||||
|
Ok(ToolResult::success(serde_json::to_value(content)?))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Write tool – writes (or overwrites) a file under the project root.
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
pub struct ResourcesWriteTool;
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct WriteArgs {
|
||||||
|
path: String,
|
||||||
|
content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for ResourcesWriteTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"resources/write"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Writes (or overwrites) a file. Requires explicit consent."
|
||||||
|
}
|
||||||
|
fn schema(&self) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": { "type": "string", "description": "Target file path (relative to project root)" },
|
||||||
|
"content": { "type": "string", "description": "File content to write" }
|
||||||
|
},
|
||||||
|
"required": ["path", "content"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
fn requires_filesystem(&self) -> Vec<String> {
|
||||||
|
vec!["file_write".to_string()]
|
||||||
|
}
|
||||||
|
async fn execute(&self, args: serde_json::Value) -> Result<ToolResult> {
|
||||||
|
let args: WriteArgs = serde_json::from_value(args)?;
|
||||||
|
let root = env::current_dir()?;
|
||||||
|
let full_path = sanitize_path(&args.path, &root)?;
|
||||||
|
// Ensure the parent directory exists
|
||||||
|
if let Some(parent) = full_path.parent() {
|
||||||
|
fs::create_dir_all(parent)?;
|
||||||
|
}
|
||||||
|
fs::write(full_path, args.content)?;
|
||||||
|
Ok(ToolResult::success(json!(null)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Delete tool – deletes a file under the project root.
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
pub struct ResourcesDeleteTool;
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct DeleteArgs {
|
||||||
|
path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for ResourcesDeleteTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"resources/delete"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Deletes a file. Requires explicit consent."
|
||||||
|
}
|
||||||
|
fn schema(&self) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": { "path": { "type": "string", "description": "File path to delete" } },
|
||||||
|
"required": ["path"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
fn requires_filesystem(&self) -> Vec<String> {
|
||||||
|
vec!["file_delete".to_string()]
|
||||||
|
}
|
||||||
|
async fn execute(&self, args: serde_json::Value) -> Result<ToolResult> {
|
||||||
|
let args: DeleteArgs = serde_json::from_value(args)?;
|
||||||
|
let root = env::current_dir()?;
|
||||||
|
let full_path = sanitize_path(&args.path, &root)?;
|
||||||
|
if full_path.is_file() {
|
||||||
|
fs::remove_file(full_path)?;
|
||||||
|
Ok(ToolResult::success(json!(null)))
|
||||||
|
} else {
|
||||||
|
Err(Error::InvalidInput("Path does not refer to a file".into()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
83
crates/owlen-core/src/tools/registry.rs
Normal file
83
crates/owlen-core/src/tools/registry.rs
Normal file
@@ -0,0 +1,83 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use crate::Result;
|
||||||
|
use anyhow::Context;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
use super::{Tool, ToolResult};
|
||||||
|
use crate::config::Config;
|
||||||
|
use crate::ui::UiController;
|
||||||
|
|
||||||
|
pub struct ToolRegistry {
|
||||||
|
tools: HashMap<String, Arc<dyn Tool>>,
|
||||||
|
config: Arc<tokio::sync::Mutex<Config>>,
|
||||||
|
ui: Arc<dyn UiController>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolRegistry {
|
||||||
|
pub fn new(config: Arc<tokio::sync::Mutex<Config>>, ui: Arc<dyn UiController>) -> Self {
|
||||||
|
Self {
|
||||||
|
tools: HashMap::new(),
|
||||||
|
config,
|
||||||
|
ui,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn register<T>(&mut self, tool: T)
|
||||||
|
where
|
||||||
|
T: Tool + 'static,
|
||||||
|
{
|
||||||
|
let tool: Arc<dyn Tool> = Arc::new(tool);
|
||||||
|
let name = tool.name().to_string();
|
||||||
|
self.tools.insert(name, tool);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
|
||||||
|
self.tools.get(name).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn all(&self) -> Vec<Arc<dyn Tool>> {
|
||||||
|
self.tools.values().cloned().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn execute(&self, name: &str, args: Value) -> Result<ToolResult> {
|
||||||
|
let tool = self
|
||||||
|
.get(name)
|
||||||
|
.with_context(|| format!("Tool not registered: {}", name))?;
|
||||||
|
|
||||||
|
let mut config = self.config.lock().await;
|
||||||
|
|
||||||
|
let is_enabled = match name {
|
||||||
|
"web_search" => config.tools.web_search.enabled,
|
||||||
|
"code_exec" => config.tools.code_exec.enabled,
|
||||||
|
_ => true, // All other tools are considered enabled by default
|
||||||
|
};
|
||||||
|
|
||||||
|
if !is_enabled {
|
||||||
|
let prompt = format!(
|
||||||
|
"Tool '{}' is disabled. Would you like to enable it for this session?",
|
||||||
|
name
|
||||||
|
);
|
||||||
|
if self.ui.confirm(&prompt).await {
|
||||||
|
// Enable the tool in the in-memory config for the current session
|
||||||
|
match name {
|
||||||
|
"web_search" => config.tools.web_search.enabled = true,
|
||||||
|
"code_exec" => config.tools.code_exec.enabled = true,
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
return Ok(ToolResult::cancelled(&format!(
|
||||||
|
"Tool '{}' execution was cancelled by the user.",
|
||||||
|
name
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tool.execute(args).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn tools(&self) -> Vec<String> {
|
||||||
|
self.tools.keys().cloned().collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
154
crates/owlen-core/src/tools/web_search.rs
Normal file
154
crates/owlen-core/src/tools/web_search.rs
Normal file
@@ -0,0 +1,154 @@
|
|||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use crate::Result;
|
||||||
|
use anyhow::Context;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
use super::{Tool, ToolResult};
|
||||||
|
use crate::consent::ConsentManager;
|
||||||
|
use crate::credentials::CredentialManager;
|
||||||
|
use crate::encryption::VaultHandle;
|
||||||
|
|
||||||
|
pub struct WebSearchTool {
|
||||||
|
consent_manager: Arc<Mutex<ConsentManager>>,
|
||||||
|
_credential_manager: Option<Arc<CredentialManager>>,
|
||||||
|
browser: duckduckgo::browser::Browser,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WebSearchTool {
|
||||||
|
pub fn new(
|
||||||
|
consent_manager: Arc<Mutex<ConsentManager>>,
|
||||||
|
credential_manager: Option<Arc<CredentialManager>>,
|
||||||
|
_vault: Option<Arc<Mutex<VaultHandle>>>,
|
||||||
|
) -> Self {
|
||||||
|
// Create a reqwest client compatible with duckduckgo crate (v0.11)
|
||||||
|
let client = reqwest_011::Client::new();
|
||||||
|
let browser = duckduckgo::browser::Browser::new(client);
|
||||||
|
|
||||||
|
Self {
|
||||||
|
consent_manager,
|
||||||
|
_credential_manager: credential_manager,
|
||||||
|
browser,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for WebSearchTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"web_search"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Search the web for information using DuckDuckGo API"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"query": {
|
||||||
|
"type": "string",
|
||||||
|
"minLength": 1,
|
||||||
|
"maxLength": 500,
|
||||||
|
"description": "Search query"
|
||||||
|
},
|
||||||
|
"max_results": {
|
||||||
|
"type": "integer",
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 10,
|
||||||
|
"default": 5,
|
||||||
|
"description": "Maximum number of results"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["query"],
|
||||||
|
"additionalProperties": false
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn requires_network(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value) -> Result<ToolResult> {
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
// Check if consent has been granted (non-blocking check)
|
||||||
|
// Consent should have been granted via TUI dialog before tool execution
|
||||||
|
{
|
||||||
|
let consent = self
|
||||||
|
.consent_manager
|
||||||
|
.lock()
|
||||||
|
.expect("Consent manager mutex poisoned");
|
||||||
|
|
||||||
|
if !consent.has_consent(self.name()) {
|
||||||
|
return Ok(ToolResult::error(
|
||||||
|
"Consent not granted for web search. This should have been handled by the TUI.",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let query = args
|
||||||
|
.get("query")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.context("Missing query parameter")?;
|
||||||
|
let max_results = args.get("max_results").and_then(Value::as_u64).unwrap_or(5) as usize;
|
||||||
|
|
||||||
|
let user_agent = duckduckgo::user_agents::get("firefox").unwrap_or(
|
||||||
|
"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:91.0) Gecko/20100101 Firefox/91.0",
|
||||||
|
);
|
||||||
|
|
||||||
|
// Detect if this is a news query - use news endpoint for better snippets
|
||||||
|
let is_news_query = query.to_lowercase().contains("news")
|
||||||
|
|| query.to_lowercase().contains("latest")
|
||||||
|
|| query.to_lowercase().contains("today")
|
||||||
|
|| query.to_lowercase().contains("recent");
|
||||||
|
|
||||||
|
let mut formatted_results = Vec::new();
|
||||||
|
|
||||||
|
if is_news_query {
|
||||||
|
// Use news endpoint which returns excerpts/snippets
|
||||||
|
let news_results = self
|
||||||
|
.browser
|
||||||
|
.news(query, "wt-wt", false, Some(max_results), user_agent)
|
||||||
|
.await
|
||||||
|
.context("DuckDuckGo news search failed")?;
|
||||||
|
|
||||||
|
for result in news_results {
|
||||||
|
formatted_results.push(json!({
|
||||||
|
"title": result.title,
|
||||||
|
"url": result.url,
|
||||||
|
"snippet": result.body, // news has body/excerpt
|
||||||
|
"source": result.source,
|
||||||
|
"date": result.date
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Use lite search for general queries (fast but no snippets)
|
||||||
|
let search_results = self
|
||||||
|
.browser
|
||||||
|
.lite_search(query, "wt-wt", Some(max_results), user_agent)
|
||||||
|
.await
|
||||||
|
.context("DuckDuckGo search failed")?;
|
||||||
|
|
||||||
|
for result in search_results {
|
||||||
|
formatted_results.push(json!({
|
||||||
|
"title": result.title,
|
||||||
|
"url": result.url,
|
||||||
|
"snippet": result.snippet
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut result = ToolResult::success(json!({
|
||||||
|
"query": query,
|
||||||
|
"results": formatted_results,
|
||||||
|
"total_found": formatted_results.len()
|
||||||
|
}));
|
||||||
|
result.duration = start.elapsed();
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
}
|
||||||
131
crates/owlen-core/src/tools/web_search_detailed.rs
Normal file
131
crates/owlen-core/src/tools/web_search_detailed.rs
Normal file
@@ -0,0 +1,131 @@
|
|||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use crate::Result;
|
||||||
|
use anyhow::Context;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
use super::{Tool, ToolResult};
|
||||||
|
use crate::consent::ConsentManager;
|
||||||
|
use crate::credentials::CredentialManager;
|
||||||
|
use crate::encryption::VaultHandle;
|
||||||
|
|
||||||
|
pub struct WebSearchDetailedTool {
|
||||||
|
consent_manager: Arc<Mutex<ConsentManager>>,
|
||||||
|
_credential_manager: Option<Arc<CredentialManager>>,
|
||||||
|
browser: duckduckgo::browser::Browser,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WebSearchDetailedTool {
|
||||||
|
pub fn new(
|
||||||
|
consent_manager: Arc<Mutex<ConsentManager>>,
|
||||||
|
credential_manager: Option<Arc<CredentialManager>>,
|
||||||
|
_vault: Option<Arc<Mutex<VaultHandle>>>,
|
||||||
|
) -> Self {
|
||||||
|
// Create a reqwest client compatible with duckduckgo crate (v0.11)
|
||||||
|
let client = reqwest_011::Client::new();
|
||||||
|
let browser = duckduckgo::browser::Browser::new(client);
|
||||||
|
|
||||||
|
Self {
|
||||||
|
consent_manager,
|
||||||
|
_credential_manager: credential_manager,
|
||||||
|
browser,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for WebSearchDetailedTool {
|
||||||
|
fn name(&self) -> &'static str {
|
||||||
|
"web_search_detailed"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &'static str {
|
||||||
|
"Search for recent articles and web content with detailed snippets and descriptions. \
|
||||||
|
Returns results with publication dates, sources, and full text excerpts. \
|
||||||
|
Best for finding recent information, articles, and detailed context about topics."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema(&self) -> Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"query": {
|
||||||
|
"type": "string",
|
||||||
|
"minLength": 1,
|
||||||
|
"maxLength": 500,
|
||||||
|
"description": "Search query"
|
||||||
|
},
|
||||||
|
"max_results": {
|
||||||
|
"type": "integer",
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 10,
|
||||||
|
"default": 5,
|
||||||
|
"description": "Maximum number of results"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["query"],
|
||||||
|
"additionalProperties": false
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn requires_network(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: Value) -> Result<ToolResult> {
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
// Check if consent has been granted (non-blocking check)
|
||||||
|
// Consent should have been granted via TUI dialog before tool execution
|
||||||
|
{
|
||||||
|
let consent = self
|
||||||
|
.consent_manager
|
||||||
|
.lock()
|
||||||
|
.expect("Consent manager mutex poisoned");
|
||||||
|
|
||||||
|
if !consent.has_consent(self.name()) {
|
||||||
|
return Ok(ToolResult::error("Consent not granted for detailed web search. This should have been handled by the TUI."));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let query = args
|
||||||
|
.get("query")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.context("Missing query parameter")?;
|
||||||
|
let max_results = args.get("max_results").and_then(Value::as_u64).unwrap_or(5) as usize;
|
||||||
|
|
||||||
|
let user_agent = duckduckgo::user_agents::get("firefox").unwrap_or(
|
||||||
|
"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:91.0) Gecko/20100101 Firefox/91.0",
|
||||||
|
);
|
||||||
|
|
||||||
|
// Use news endpoint which provides detailed results with full snippets
|
||||||
|
// Even for non-news queries, this often returns recent articles and content with good descriptions
|
||||||
|
let news_results = self
|
||||||
|
.browser
|
||||||
|
.news(query, "wt-wt", false, Some(max_results), user_agent)
|
||||||
|
.await
|
||||||
|
.context("DuckDuckGo detailed search failed")?;
|
||||||
|
|
||||||
|
let mut formatted_results = Vec::new();
|
||||||
|
for result in news_results {
|
||||||
|
formatted_results.push(json!({
|
||||||
|
"title": result.title,
|
||||||
|
"url": result.url,
|
||||||
|
"snippet": result.body, // news endpoint includes full excerpts
|
||||||
|
"source": result.source,
|
||||||
|
"date": result.date
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut result = ToolResult::success(json!({
|
||||||
|
"query": query,
|
||||||
|
"results": formatted_results,
|
||||||
|
"total_found": formatted_results.len()
|
||||||
|
}));
|
||||||
|
result.duration = start.elapsed();
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -18,6 +18,9 @@ pub struct Message {
|
|||||||
pub metadata: HashMap<String, serde_json::Value>,
|
pub metadata: HashMap<String, serde_json::Value>,
|
||||||
/// Timestamp when the message was created
|
/// Timestamp when the message was created
|
||||||
pub timestamp: std::time::SystemTime,
|
pub timestamp: std::time::SystemTime,
|
||||||
|
/// Tool calls requested by the assistant
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tool_calls: Option<Vec<ToolCall>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Role of a message sender
|
/// Role of a message sender
|
||||||
@@ -30,6 +33,19 @@ pub enum Role {
|
|||||||
Assistant,
|
Assistant,
|
||||||
/// System message (prompts, context, etc.)
|
/// System message (prompts, context, etc.)
|
||||||
System,
|
System,
|
||||||
|
/// Tool response message
|
||||||
|
Tool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A tool call requested by the assistant
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub struct ToolCall {
|
||||||
|
/// Unique identifier for this tool call
|
||||||
|
pub id: String,
|
||||||
|
/// Name of the tool to call
|
||||||
|
pub name: String,
|
||||||
|
/// Arguments for the tool (JSON object)
|
||||||
|
pub arguments: serde_json::Value,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl fmt::Display for Role {
|
impl fmt::Display for Role {
|
||||||
@@ -38,6 +54,7 @@ impl fmt::Display for Role {
|
|||||||
Role::User => "user",
|
Role::User => "user",
|
||||||
Role::Assistant => "assistant",
|
Role::Assistant => "assistant",
|
||||||
Role::System => "system",
|
Role::System => "system",
|
||||||
|
Role::Tool => "tool",
|
||||||
};
|
};
|
||||||
f.write_str(label)
|
f.write_str(label)
|
||||||
}
|
}
|
||||||
@@ -72,6 +89,9 @@ pub struct ChatRequest {
|
|||||||
pub messages: Vec<Message>,
|
pub messages: Vec<Message>,
|
||||||
/// Optional parameters for the request
|
/// Optional parameters for the request
|
||||||
pub parameters: ChatParameters,
|
pub parameters: ChatParameters,
|
||||||
|
/// Optional tools available for the model to use
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tools: Option<Vec<crate::mcp::McpToolDescriptor>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Parameters for chat completion
|
/// Parameters for chat completion
|
||||||
@@ -133,6 +153,9 @@ pub struct ModelInfo {
|
|||||||
pub context_window: Option<u32>,
|
pub context_window: Option<u32>,
|
||||||
/// Additional capabilities
|
/// Additional capabilities
|
||||||
pub capabilities: Vec<String>,
|
pub capabilities: Vec<String>,
|
||||||
|
/// Whether this model supports tool/function calling
|
||||||
|
#[serde(default)]
|
||||||
|
pub supports_tools: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Message {
|
impl Message {
|
||||||
@@ -144,6 +167,7 @@ impl Message {
|
|||||||
content,
|
content,
|
||||||
metadata: HashMap::new(),
|
metadata: HashMap::new(),
|
||||||
timestamp: std::time::SystemTime::now(),
|
timestamp: std::time::SystemTime::now(),
|
||||||
|
tool_calls: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -161,6 +185,24 @@ impl Message {
|
|||||||
pub fn system(content: String) -> Self {
|
pub fn system(content: String) -> Self {
|
||||||
Self::new(Role::System, content)
|
Self::new(Role::System, content)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Create a tool response message
|
||||||
|
pub fn tool(tool_call_id: String, content: String) -> Self {
|
||||||
|
let mut msg = Self::new(Role::Tool, content);
|
||||||
|
msg.metadata.insert(
|
||||||
|
"tool_call_id".to_string(),
|
||||||
|
serde_json::Value::String(tool_call_id),
|
||||||
|
);
|
||||||
|
msg
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if this message has tool calls
|
||||||
|
pub fn has_tool_calls(&self) -> bool {
|
||||||
|
self.tool_calls
|
||||||
|
.as_ref()
|
||||||
|
.map(|tc| !tc.is_empty())
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Conversation {
|
impl Conversation {
|
||||||
|
|||||||
@@ -351,14 +351,52 @@ pub fn find_prev_word_boundary(line: &str, col: usize) -> Option<usize> {
|
|||||||
Some(pos)
|
Some(pos)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
use crate::theme::Theme;
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use std::io::stdout;
|
||||||
|
|
||||||
|
pub fn show_mouse_cursor() {
|
||||||
|
let mut stdout = stdout();
|
||||||
|
crossterm::execute!(stdout, crossterm::cursor::Show).ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn hide_mouse_cursor() {
|
||||||
|
let mut stdout = stdout();
|
||||||
|
crossterm::execute!(stdout, crossterm::cursor::Hide).ok();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn apply_theme_to_string(s: &str, _theme: &Theme) -> String {
|
||||||
|
// This is a placeholder. In a real implementation, you'd parse the string
|
||||||
|
// and apply colors based on syntax or other rules.
|
||||||
|
s.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A trait for abstracting UI interactions like confirmations.
|
||||||
|
#[async_trait]
|
||||||
|
pub trait UiController: Send + Sync {
|
||||||
|
async fn confirm(&self, prompt: &str) -> bool;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A no-op UI controller for non-interactive contexts.
|
||||||
|
pub struct NoOpUiController;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl UiController for NoOpUiController {
|
||||||
|
async fn confirm(&self, _prompt: &str) -> bool {
|
||||||
|
false // Always decline in non-interactive mode
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_auto_scroll() {
|
fn test_auto_scroll() {
|
||||||
let mut scroll = AutoScroll::default();
|
let mut scroll = AutoScroll {
|
||||||
scroll.content_len = 100;
|
content_len: 100,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
// Test on_viewport with stick_to_bottom
|
// Test on_viewport with stick_to_bottom
|
||||||
scroll.on_viewport(10);
|
scroll.on_viewport(10);
|
||||||
|
|||||||
108
crates/owlen-core/src/validation.rs
Normal file
108
crates/owlen-core/src/validation.rs
Normal file
@@ -0,0 +1,108 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use anyhow::{Context, Result};
|
||||||
|
use jsonschema::{JSONSchema, ValidationError};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
pub struct SchemaValidator {
|
||||||
|
schemas: HashMap<String, JSONSchema>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for SchemaValidator {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SchemaValidator {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
schemas: HashMap::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn register_schema(&mut self, tool_name: &str, schema: Value) -> Result<()> {
|
||||||
|
let compiled = JSONSchema::compile(&schema)
|
||||||
|
.map_err(|e| anyhow::anyhow!("Invalid schema for {}: {}", tool_name, e))?;
|
||||||
|
|
||||||
|
self.schemas.insert(tool_name.to_string(), compiled);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn validate(&self, tool_name: &str, input: &Value) -> Result<()> {
|
||||||
|
let schema = self
|
||||||
|
.schemas
|
||||||
|
.get(tool_name)
|
||||||
|
.with_context(|| format!("No schema registered for tool: {}", tool_name))?;
|
||||||
|
|
||||||
|
if let Err(errors) = schema.validate(input) {
|
||||||
|
let error_messages: Vec<String> = errors.map(format_validation_error).collect();
|
||||||
|
|
||||||
|
return Err(anyhow::anyhow!(
|
||||||
|
"Input validation failed for {}: {}",
|
||||||
|
tool_name,
|
||||||
|
error_messages.join(", ")
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_validation_error(error: ValidationError) -> String {
|
||||||
|
format!("Validation error at {}: {}", error.instance_path, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_builtin_schemas() -> HashMap<String, Value> {
|
||||||
|
let mut schemas = HashMap::new();
|
||||||
|
|
||||||
|
schemas.insert(
|
||||||
|
"web_search".to_string(),
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"query": {
|
||||||
|
"type": "string",
|
||||||
|
"minLength": 1,
|
||||||
|
"maxLength": 500
|
||||||
|
},
|
||||||
|
"max_results": {
|
||||||
|
"type": "integer",
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 10,
|
||||||
|
"default": 5
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["query"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
schemas.insert(
|
||||||
|
"code_exec".to_string(),
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"language": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["python", "javascript", "bash", "rust"]
|
||||||
|
},
|
||||||
|
"code": {
|
||||||
|
"type": "string",
|
||||||
|
"minLength": 1,
|
||||||
|
"maxLength": 10000
|
||||||
|
},
|
||||||
|
"timeout": {
|
||||||
|
"type": "integer",
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 300,
|
||||||
|
"default": 30
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["language", "code"],
|
||||||
|
"additionalProperties": false
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
schemas
|
||||||
|
}
|
||||||
99
crates/owlen-core/tests/consent_scope.rs
Normal file
99
crates/owlen-core/tests/consent_scope.rs
Normal file
@@ -0,0 +1,99 @@
|
|||||||
|
use owlen_core::consent::{ConsentManager, ConsentScope};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_consent_scopes() {
|
||||||
|
let mut manager = ConsentManager::new();
|
||||||
|
|
||||||
|
// Test session consent
|
||||||
|
manager.grant_consent_with_scope(
|
||||||
|
"test_tool",
|
||||||
|
vec!["data".to_string()],
|
||||||
|
vec!["https://example.com".to_string()],
|
||||||
|
ConsentScope::Session,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(manager.has_consent("test_tool"));
|
||||||
|
|
||||||
|
// Clear session consent and verify it's gone
|
||||||
|
manager.clear_session_consent();
|
||||||
|
assert!(!manager.has_consent("test_tool"));
|
||||||
|
|
||||||
|
// Test permanent consent survives session clear
|
||||||
|
manager.grant_consent_with_scope(
|
||||||
|
"test_tool_permanent",
|
||||||
|
vec!["data".to_string()],
|
||||||
|
vec!["https://example.com".to_string()],
|
||||||
|
ConsentScope::Permanent,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(manager.has_consent("test_tool_permanent"));
|
||||||
|
manager.clear_session_consent();
|
||||||
|
assert!(manager.has_consent("test_tool_permanent"));
|
||||||
|
|
||||||
|
// Verify revoke works for permanent consent
|
||||||
|
manager.revoke_consent("test_tool_permanent");
|
||||||
|
assert!(!manager.has_consent("test_tool_permanent"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_pending_requests_prevents_duplicates() {
|
||||||
|
let mut manager = ConsentManager::new();
|
||||||
|
|
||||||
|
// Simulate concurrent consent requests by checking pending state
|
||||||
|
// In real usage, multiple threads would call request_consent simultaneously
|
||||||
|
|
||||||
|
// First, verify a tool has no consent
|
||||||
|
assert!(!manager.has_consent("web_search"));
|
||||||
|
|
||||||
|
// The pending_requests map is private, but we can test the behavior
|
||||||
|
// by checking that consent checks work correctly
|
||||||
|
assert!(manager.check_consent_needed("web_search").is_some());
|
||||||
|
|
||||||
|
// Grant session consent
|
||||||
|
manager.grant_consent_with_scope(
|
||||||
|
"web_search",
|
||||||
|
vec!["search queries".to_string()],
|
||||||
|
vec!["https://api.search.com".to_string()],
|
||||||
|
ConsentScope::Session,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Now it should have consent
|
||||||
|
assert!(manager.has_consent("web_search"));
|
||||||
|
assert!(manager.check_consent_needed("web_search").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_consent_record_separation() {
|
||||||
|
let mut manager = ConsentManager::new();
|
||||||
|
|
||||||
|
// Add permanent consent
|
||||||
|
manager.grant_consent_with_scope(
|
||||||
|
"perm_tool",
|
||||||
|
vec!["data".to_string()],
|
||||||
|
vec!["https://perm.com".to_string()],
|
||||||
|
ConsentScope::Permanent,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Add session consent
|
||||||
|
manager.grant_consent_with_scope(
|
||||||
|
"session_tool",
|
||||||
|
vec!["data".to_string()],
|
||||||
|
vec!["https://session.com".to_string()],
|
||||||
|
ConsentScope::Session,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Both should have consent
|
||||||
|
assert!(manager.has_consent("perm_tool"));
|
||||||
|
assert!(manager.has_consent("session_tool"));
|
||||||
|
|
||||||
|
// Clear session consent
|
||||||
|
manager.clear_session_consent();
|
||||||
|
|
||||||
|
// Only permanent should remain
|
||||||
|
assert!(manager.has_consent("perm_tool"));
|
||||||
|
assert!(!manager.has_consent("session_tool"));
|
||||||
|
|
||||||
|
// Clear all
|
||||||
|
manager.clear_all_consent();
|
||||||
|
assert!(!manager.has_consent("perm_tool"));
|
||||||
|
}
|
||||||
53
crates/owlen-core/tests/file_server.rs
Normal file
53
crates/owlen-core/tests/file_server.rs
Normal file
@@ -0,0 +1,53 @@
|
|||||||
|
use owlen_core::mcp::client::McpClient;
|
||||||
|
use owlen_core::mcp::remote_client::RemoteMcpClient;
|
||||||
|
use owlen_core::mcp::McpToolCall;
|
||||||
|
use std::fs::File;
|
||||||
|
use std::io::Write;
|
||||||
|
use tempfile::tempdir;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn remote_file_server_read_and_list() {
|
||||||
|
// Create temporary directory with a file
|
||||||
|
let dir = tempdir().expect("tempdir failed");
|
||||||
|
let file_path = dir.path().join("hello.txt");
|
||||||
|
let mut file = File::create(&file_path).expect("create file");
|
||||||
|
writeln!(file, "world").expect("write file");
|
||||||
|
|
||||||
|
// Change current directory for the test process so the server sees the temp dir as its root
|
||||||
|
std::env::set_current_dir(dir.path()).expect("set cwd");
|
||||||
|
|
||||||
|
// Ensure the MCP server binary is built.
|
||||||
|
// Build the MCP server binary using the workspace manifest.
|
||||||
|
let manifest_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||||
|
.join("../..")
|
||||||
|
.join("Cargo.toml");
|
||||||
|
let build_status = std::process::Command::new("cargo")
|
||||||
|
.args(&["build", "-p", "owlen-mcp-server", "--manifest-path"])
|
||||||
|
.arg(manifest_path)
|
||||||
|
.status()
|
||||||
|
.expect("failed to run cargo build for MCP server");
|
||||||
|
assert!(build_status.success(), "MCP server build failed");
|
||||||
|
|
||||||
|
// Spawn remote client after the cwd is set and binary built
|
||||||
|
let client = RemoteMcpClient::new().expect("remote client init");
|
||||||
|
|
||||||
|
// Read file via MCP
|
||||||
|
let call = McpToolCall {
|
||||||
|
name: "resources/get".to_string(),
|
||||||
|
arguments: serde_json::json!({"path": "hello.txt"}),
|
||||||
|
};
|
||||||
|
let resp = client.call_tool(call).await.expect("call_tool");
|
||||||
|
let content: String = serde_json::from_value(resp.output).expect("parse output");
|
||||||
|
assert!(content.trim().ends_with("world"));
|
||||||
|
|
||||||
|
// List directory via MCP
|
||||||
|
let list_call = McpToolCall {
|
||||||
|
name: "resources/list".to_string(),
|
||||||
|
arguments: serde_json::json!({"path": "."}),
|
||||||
|
};
|
||||||
|
let list_resp = client.call_tool(list_call).await.expect("list_tool");
|
||||||
|
let entries: Vec<String> = serde_json::from_value(list_resp.output).expect("parse list");
|
||||||
|
assert!(entries.contains(&"hello.txt".to_string()));
|
||||||
|
|
||||||
|
// Cleanup handled by tempdir
|
||||||
|
}
|
||||||
68
crates/owlen-core/tests/file_write.rs
Normal file
68
crates/owlen-core/tests/file_write.rs
Normal file
@@ -0,0 +1,68 @@
|
|||||||
|
use owlen_core::mcp::client::McpClient;
|
||||||
|
use owlen_core::mcp::remote_client::RemoteMcpClient;
|
||||||
|
use owlen_core::mcp::McpToolCall;
|
||||||
|
use tempfile::tempdir;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn remote_write_and_delete() {
|
||||||
|
// Build the server binary first
|
||||||
|
let status = std::process::Command::new("cargo")
|
||||||
|
.args(&["build", "-p", "owlen-mcp-server"])
|
||||||
|
.status()
|
||||||
|
.expect("failed to build MCP server");
|
||||||
|
assert!(status.success());
|
||||||
|
|
||||||
|
// Use a temp dir as project root
|
||||||
|
let dir = tempdir().expect("tempdir");
|
||||||
|
std::env::set_current_dir(dir.path()).expect("set cwd");
|
||||||
|
|
||||||
|
let client = RemoteMcpClient::new().expect("client init");
|
||||||
|
|
||||||
|
// Write a file via MCP
|
||||||
|
let write_call = McpToolCall {
|
||||||
|
name: "resources/write".to_string(),
|
||||||
|
arguments: serde_json::json!({ "path": "test.txt", "content": "hello" }),
|
||||||
|
};
|
||||||
|
client.call_tool(write_call).await.expect("write tool");
|
||||||
|
|
||||||
|
// Verify content via local read (fallback check)
|
||||||
|
let content = std::fs::read_to_string(dir.path().join("test.txt")).expect("read back");
|
||||||
|
assert_eq!(content, "hello");
|
||||||
|
|
||||||
|
// Delete the file via MCP
|
||||||
|
let del_call = McpToolCall {
|
||||||
|
name: "resources/delete".to_string(),
|
||||||
|
arguments: serde_json::json!({ "path": "test.txt" }),
|
||||||
|
};
|
||||||
|
client.call_tool(del_call).await.expect("delete tool");
|
||||||
|
assert!(!dir.path().join("test.txt").exists());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn write_outside_root_is_rejected() {
|
||||||
|
// Build server (already built in previous test, but ensure it exists)
|
||||||
|
let status = std::process::Command::new("cargo")
|
||||||
|
.args(&["build", "-p", "owlen-mcp-server"])
|
||||||
|
.status()
|
||||||
|
.expect("failed to build MCP server");
|
||||||
|
assert!(status.success());
|
||||||
|
|
||||||
|
// Set cwd to a fresh temp dir
|
||||||
|
let dir = tempdir().expect("tempdir");
|
||||||
|
std::env::set_current_dir(dir.path()).expect("set cwd");
|
||||||
|
let client = RemoteMcpClient::new().expect("client init");
|
||||||
|
|
||||||
|
// Attempt to write outside the root using "../evil.txt"
|
||||||
|
let call = McpToolCall {
|
||||||
|
name: "resources/write".to_string(),
|
||||||
|
arguments: serde_json::json!({ "path": "../evil.txt", "content": "bad" }),
|
||||||
|
};
|
||||||
|
let err = client.call_tool(call).await.unwrap_err();
|
||||||
|
// The server returns a Network error with path traversal message
|
||||||
|
let err_str = format!("{err}");
|
||||||
|
assert!(
|
||||||
|
err_str.contains("path traversal") || err_str.contains("Path traversal"),
|
||||||
|
"Expected path traversal error, got: {}",
|
||||||
|
err_str
|
||||||
|
);
|
||||||
|
}
|
||||||
12
crates/owlen-mcp-client/Cargo.toml
Normal file
12
crates/owlen-mcp-client/Cargo.toml
Normal file
@@ -0,0 +1,12 @@
|
|||||||
|
[package]
|
||||||
|
name = "owlen-mcp-client"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
description = "Dedicated MCP client library for Owlen, exposing remote MCP server communication"
|
||||||
|
license = "AGPL-3.0"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
owlen-core = { path = "../owlen-core" }
|
||||||
|
|
||||||
|
[features]
|
||||||
|
default = []
|
||||||
19
crates/owlen-mcp-client/src/lib.rs
Normal file
19
crates/owlen-mcp-client/src/lib.rs
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
//! Owlen MCP client library.
|
||||||
|
//!
|
||||||
|
//! This crate provides a thin façade over the remote MCP client implementation
|
||||||
|
//! inside `owlen-core`. It re‑exports the most useful types so downstream
|
||||||
|
//! crates can depend only on `owlen-mcp-client` without pulling in the entire
|
||||||
|
//! core crate internals.
|
||||||
|
|
||||||
|
pub use owlen_core::mcp::remote_client::RemoteMcpClient;
|
||||||
|
pub use owlen_core::mcp::{McpClient, McpToolCall, McpToolDescriptor, McpToolResponse};
|
||||||
|
|
||||||
|
// Re‑export the Provider implementation so the client can also be used as an
|
||||||
|
// LLM provider when the remote MCP server hosts a language‑model tool (e.g.
|
||||||
|
// `generate_text`).
|
||||||
|
// Re‑export the core Provider trait so that the MCP client can also be used as an LLM provider.
|
||||||
|
pub use owlen_core::provider::Provider as McpProvider;
|
||||||
|
|
||||||
|
// Note: The `RemoteMcpClient` type provides its own `new` constructor in the core
|
||||||
|
// crate. Users can call `RemoteMcpClient::new()` directly. No additional wrapper
|
||||||
|
// is needed here.
|
||||||
20
crates/owlen-mcp-llm-server/Cargo.toml
Normal file
20
crates/owlen-mcp-llm-server/Cargo.toml
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
[package]
|
||||||
|
name = "owlen-mcp-llm-server"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
owlen-core = { path = "../owlen-core" }
|
||||||
|
owlen-ollama = { path = "../owlen-ollama" }
|
||||||
|
tokio = { version = "1.0", features = ["full"] }
|
||||||
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
|
serde_json = "1.0"
|
||||||
|
anyhow = "1.0"
|
||||||
|
tokio-stream = "0.1"
|
||||||
|
|
||||||
|
[lib]
|
||||||
|
path = "src/lib.rs"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "owlen-mcp-llm-server"
|
||||||
|
path = "src/lib.rs"
|
||||||
498
crates/owlen-mcp-llm-server/src/lib.rs
Normal file
498
crates/owlen-mcp-llm-server/src/lib.rs
Normal file
@@ -0,0 +1,498 @@
|
|||||||
|
#![allow(
|
||||||
|
unused_imports,
|
||||||
|
unused_variables,
|
||||||
|
dead_code,
|
||||||
|
clippy::unnecessary_cast,
|
||||||
|
clippy::manual_flatten,
|
||||||
|
clippy::empty_line_after_outer_attr
|
||||||
|
)]
|
||||||
|
|
||||||
|
use owlen_core::mcp::protocol::{
|
||||||
|
methods, ErrorCode, InitializeParams, InitializeResult, RequestId, RpcError, RpcErrorResponse,
|
||||||
|
RpcNotification, RpcRequest, RpcResponse, ServerCapabilities, ServerInfo, PROTOCOL_VERSION,
|
||||||
|
};
|
||||||
|
use owlen_core::mcp::{McpToolCall, McpToolDescriptor, McpToolResponse};
|
||||||
|
use owlen_core::types::{ChatParameters, ChatRequest, Message};
|
||||||
|
use owlen_core::Provider;
|
||||||
|
use owlen_ollama::OllamaProvider;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::env;
|
||||||
|
use tokio::io::{self, AsyncBufReadExt, AsyncWriteExt};
|
||||||
|
use tokio_stream::StreamExt;
|
||||||
|
|
||||||
|
// Suppress warnings are handled by the crate-level attribute at the top.
|
||||||
|
|
||||||
|
/// Arguments for the generate_text tool
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct GenerateTextArgs {
|
||||||
|
messages: Vec<Message>,
|
||||||
|
temperature: Option<f32>,
|
||||||
|
max_tokens: Option<u32>,
|
||||||
|
model: String,
|
||||||
|
stream: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Simple tool descriptor for generate_text
|
||||||
|
fn generate_text_descriptor() -> McpToolDescriptor {
|
||||||
|
McpToolDescriptor {
|
||||||
|
name: "generate_text".to_string(),
|
||||||
|
description: "Generate text using Ollama LLM. Each message must have 'role' (user/assistant/system) and 'content' (string) fields.".to_string(),
|
||||||
|
input_schema: json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"messages": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"role": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["user", "assistant", "system"],
|
||||||
|
"description": "The role of the message sender"
|
||||||
|
},
|
||||||
|
"content": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "The message content"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["role", "content"]
|
||||||
|
},
|
||||||
|
"description": "Array of message objects with role and content"
|
||||||
|
},
|
||||||
|
"temperature": {"type": ["number", "null"], "description": "Sampling temperature (0.0-2.0)"},
|
||||||
|
"max_tokens": {"type": ["integer", "null"], "description": "Maximum tokens to generate"},
|
||||||
|
"model": {"type": "string", "description": "Model name (e.g., llama3.2:latest)"},
|
||||||
|
"stream": {"type": "boolean", "description": "Whether to stream the response"}
|
||||||
|
},
|
||||||
|
"required": ["messages", "model", "stream"]
|
||||||
|
}),
|
||||||
|
requires_network: true,
|
||||||
|
requires_filesystem: vec![],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tool descriptor for resources/get (read file)
|
||||||
|
fn resources_get_descriptor() -> McpToolDescriptor {
|
||||||
|
McpToolDescriptor {
|
||||||
|
name: "resources/get".to_string(),
|
||||||
|
description: "Read and return the TEXT CONTENTS of a single FILE. Use this to read the contents of code files, config files, or text documents. Do NOT use for directories.".to_string(),
|
||||||
|
input_schema: json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Path to the FILE (not directory) to read"}
|
||||||
|
},
|
||||||
|
"required": ["path"]
|
||||||
|
}),
|
||||||
|
requires_network: false,
|
||||||
|
requires_filesystem: vec!["read".to_string()],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tool descriptor for resources/list (list directory)
|
||||||
|
fn resources_list_descriptor() -> McpToolDescriptor {
|
||||||
|
McpToolDescriptor {
|
||||||
|
name: "resources/list".to_string(),
|
||||||
|
description: "List the NAMES of all files and directories in a directory. Use this to see what files exist in a folder, or to list directory contents. Returns an array of file/directory names.".to_string(),
|
||||||
|
input_schema: json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "Path to the DIRECTORY to list (use '.' for current directory)"}
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
requires_network: false,
|
||||||
|
requires_filesystem: vec!["read".to_string()],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_generate_text(args: GenerateTextArgs) -> Result<String, RpcError> {
|
||||||
|
// Create provider with default local Ollama URL
|
||||||
|
let provider = OllamaProvider::new("http://localhost:11434")
|
||||||
|
.map_err(|e| RpcError::internal_error(format!("Failed to init OllamaProvider: {}", e)))?;
|
||||||
|
|
||||||
|
let parameters = ChatParameters {
|
||||||
|
temperature: args.temperature,
|
||||||
|
max_tokens: args.max_tokens.map(|v| v as u32),
|
||||||
|
stream: args.stream,
|
||||||
|
extra: HashMap::new(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let request = ChatRequest {
|
||||||
|
model: args.model,
|
||||||
|
messages: args.messages,
|
||||||
|
parameters,
|
||||||
|
tools: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Use streaming API and collect output
|
||||||
|
let mut stream = provider
|
||||||
|
.chat_stream(request)
|
||||||
|
.await
|
||||||
|
.map_err(|e| RpcError::internal_error(format!("Chat request failed: {}", e)))?;
|
||||||
|
let mut content = String::new();
|
||||||
|
while let Some(chunk) = stream.next().await {
|
||||||
|
match chunk {
|
||||||
|
Ok(resp) => {
|
||||||
|
content.push_str(&resp.message.content);
|
||||||
|
if resp.is_final {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
return Err(RpcError::internal_error(format!("Stream error: {}", e)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(content)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_request(req: &RpcRequest) -> Result<Value, RpcError> {
|
||||||
|
match req.method.as_str() {
|
||||||
|
methods::INITIALIZE => {
|
||||||
|
let params = req
|
||||||
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| RpcError::invalid_params("Missing params for initialize"))?;
|
||||||
|
let init: InitializeParams = serde_json::from_value(params.clone())
|
||||||
|
.map_err(|e| RpcError::invalid_params(format!("Invalid init params: {}", e)))?;
|
||||||
|
if !init.protocol_version.eq(PROTOCOL_VERSION) {
|
||||||
|
return Err(RpcError::new(
|
||||||
|
ErrorCode::INVALID_REQUEST,
|
||||||
|
format!(
|
||||||
|
"Incompatible protocol version. Client: {}, Server: {}",
|
||||||
|
init.protocol_version, PROTOCOL_VERSION
|
||||||
|
),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let result = InitializeResult {
|
||||||
|
protocol_version: PROTOCOL_VERSION.to_string(),
|
||||||
|
server_info: ServerInfo {
|
||||||
|
name: "owlen-mcp-llm-server".to_string(),
|
||||||
|
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||||
|
},
|
||||||
|
capabilities: ServerCapabilities {
|
||||||
|
supports_tools: Some(true),
|
||||||
|
supports_resources: Some(false),
|
||||||
|
supports_streaming: Some(true),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
Ok(serde_json::to_value(result).unwrap())
|
||||||
|
}
|
||||||
|
methods::TOOLS_LIST => {
|
||||||
|
let tools = vec![
|
||||||
|
generate_text_descriptor(),
|
||||||
|
resources_get_descriptor(),
|
||||||
|
resources_list_descriptor(),
|
||||||
|
];
|
||||||
|
Ok(json!(tools))
|
||||||
|
}
|
||||||
|
// New method to list available Ollama models via the provider.
|
||||||
|
methods::MODELS_LIST => {
|
||||||
|
// Reuse the provider instance for model listing.
|
||||||
|
let provider = OllamaProvider::new("http://localhost:11434").map_err(|e| {
|
||||||
|
RpcError::internal_error(format!("Failed to init OllamaProvider: {}", e))
|
||||||
|
})?;
|
||||||
|
let models = provider
|
||||||
|
.list_models()
|
||||||
|
.await
|
||||||
|
.map_err(|e| RpcError::internal_error(format!("Failed to list models: {}", e)))?;
|
||||||
|
Ok(serde_json::to_value(models).unwrap())
|
||||||
|
}
|
||||||
|
methods::TOOLS_CALL => {
|
||||||
|
// For streaming we will send incremental notifications directly from here.
|
||||||
|
// The caller (main loop) will handle writing the final response.
|
||||||
|
Err(RpcError::internal_error(
|
||||||
|
"TOOLS_CALL should be handled in main loop for streaming",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
_ => Err(RpcError::method_not_found(&req.method)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
let root = env::current_dir()?; // not used but kept for parity
|
||||||
|
let mut stdin = io::BufReader::new(io::stdin());
|
||||||
|
let mut stdout = io::stdout();
|
||||||
|
loop {
|
||||||
|
let mut line = String::new();
|
||||||
|
match stdin.read_line(&mut line).await {
|
||||||
|
Ok(0) => break,
|
||||||
|
Ok(_) => {
|
||||||
|
let req: RpcRequest = match serde_json::from_str(&line) {
|
||||||
|
Ok(r) => r,
|
||||||
|
Err(e) => {
|
||||||
|
let err = RpcErrorResponse::new(
|
||||||
|
RequestId::Number(0),
|
||||||
|
RpcError::parse_error(format!("Parse error: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let id = req.id.clone();
|
||||||
|
// Streaming tool calls (generate_text) are handled specially to emit incremental notifications.
|
||||||
|
if req.method == methods::TOOLS_CALL {
|
||||||
|
// Parse the tool call
|
||||||
|
let params = match &req.params {
|
||||||
|
Some(p) => p,
|
||||||
|
None => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::invalid_params("Missing params for tool call"),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let call: McpToolCall = match serde_json::from_value(params.clone()) {
|
||||||
|
Ok(c) => c,
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::invalid_params(format!("Invalid tool call: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
// Dispatch based on the requested tool name.
|
||||||
|
// Handle resources tools manually.
|
||||||
|
if call.name.starts_with("resources/get") {
|
||||||
|
let path = call
|
||||||
|
.arguments
|
||||||
|
.get("path")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("");
|
||||||
|
match std::fs::read_to_string(path) {
|
||||||
|
Ok(content) => {
|
||||||
|
let response = McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: true,
|
||||||
|
output: json!(content),
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
duration_ms: 0,
|
||||||
|
};
|
||||||
|
let final_resp = RpcResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
serde_json::to_value(response).unwrap(),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&final_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::internal_error(format!("Failed to read file: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if call.name.starts_with("resources/list") {
|
||||||
|
let path = call
|
||||||
|
.arguments
|
||||||
|
.get("path")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or(".");
|
||||||
|
match std::fs::read_dir(path) {
|
||||||
|
Ok(entries) => {
|
||||||
|
let mut names = Vec::new();
|
||||||
|
for entry in entries.flatten() {
|
||||||
|
if let Some(name) = entry.file_name().to_str() {
|
||||||
|
names.push(name.to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let response = McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: true,
|
||||||
|
output: json!(names),
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
duration_ms: 0,
|
||||||
|
};
|
||||||
|
let final_resp = RpcResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
serde_json::to_value(response).unwrap(),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&final_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::internal_error(format!("Failed to list dir: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Expect generate_text tool for the remaining path.
|
||||||
|
if call.name != "generate_text" {
|
||||||
|
let err_resp =
|
||||||
|
RpcErrorResponse::new(id.clone(), RpcError::tool_not_found(&call.name));
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let args: GenerateTextArgs =
|
||||||
|
match serde_json::from_value(call.arguments.clone()) {
|
||||||
|
Ok(a) => a,
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::invalid_params(format!("Invalid arguments: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Initialize Ollama provider and start streaming
|
||||||
|
let provider = match OllamaProvider::new("http://localhost:11434") {
|
||||||
|
Ok(p) => p,
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::internal_error(format!(
|
||||||
|
"Failed to init OllamaProvider: {}",
|
||||||
|
e
|
||||||
|
)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let parameters = ChatParameters {
|
||||||
|
temperature: args.temperature,
|
||||||
|
max_tokens: args.max_tokens.map(|v| v as u32),
|
||||||
|
stream: true,
|
||||||
|
extra: HashMap::new(),
|
||||||
|
};
|
||||||
|
let request = ChatRequest {
|
||||||
|
model: args.model,
|
||||||
|
messages: args.messages,
|
||||||
|
parameters,
|
||||||
|
tools: None,
|
||||||
|
};
|
||||||
|
let mut stream = match provider.chat_stream(request).await {
|
||||||
|
Ok(s) => s,
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::internal_error(format!("Chat request failed: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
// Accumulate full content while sending incremental progress notifications
|
||||||
|
let mut final_content = String::new();
|
||||||
|
while let Some(chunk) = stream.next().await {
|
||||||
|
match chunk {
|
||||||
|
Ok(resp) => {
|
||||||
|
// Append chunk to the final content buffer
|
||||||
|
final_content.push_str(&resp.message.content);
|
||||||
|
// Emit a progress notification for the UI
|
||||||
|
let notif = RpcNotification::new(
|
||||||
|
"tools/call/progress",
|
||||||
|
Some(json!({ "content": resp.message.content })),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(¬if)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
if resp.is_final {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
id.clone(),
|
||||||
|
RpcError::internal_error(format!("Stream error: {}", e)),
|
||||||
|
);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// After streaming, send the final tool response containing the full content
|
||||||
|
let final_output = final_content.clone();
|
||||||
|
let response = McpToolResponse {
|
||||||
|
name: call.name,
|
||||||
|
success: true,
|
||||||
|
output: json!(final_output),
|
||||||
|
metadata: HashMap::new(),
|
||||||
|
duration_ms: 0,
|
||||||
|
};
|
||||||
|
let final_resp =
|
||||||
|
RpcResponse::new(id.clone(), serde_json::to_value(response).unwrap());
|
||||||
|
let s = serde_json::to_string(&final_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
// Non‑streaming requests are handled by the generic handler
|
||||||
|
match handle_request(&req).await {
|
||||||
|
Ok(res) => {
|
||||||
|
let resp = RpcResponse::new(id, res);
|
||||||
|
let s = serde_json::to_string(&resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(id, err);
|
||||||
|
let s = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(s.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("Read error: {}", e);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
12
crates/owlen-mcp-server/Cargo.toml
Normal file
12
crates/owlen-mcp-server/Cargo.toml
Normal file
@@ -0,0 +1,12 @@
|
|||||||
|
[package]
|
||||||
|
name = "owlen-mcp-server"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
tokio = { version = "1.0", features = ["full"] }
|
||||||
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
|
serde_json = "1.0"
|
||||||
|
anyhow = "1.0"
|
||||||
|
path-clean = "1.0"
|
||||||
|
owlen-core = { path = "../owlen-core" }
|
||||||
246
crates/owlen-mcp-server/src/main.rs
Normal file
246
crates/owlen-mcp-server/src/main.rs
Normal file
@@ -0,0 +1,246 @@
|
|||||||
|
use owlen_core::mcp::protocol::{
|
||||||
|
is_compatible, ErrorCode, InitializeParams, InitializeResult, RequestId, RpcError,
|
||||||
|
RpcErrorResponse, RpcRequest, RpcResponse, ServerCapabilities, ServerInfo, PROTOCOL_VERSION,
|
||||||
|
};
|
||||||
|
use path_clean::PathClean;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use std::env;
|
||||||
|
use std::fs;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use tokio::io::{self, AsyncBufReadExt, AsyncWriteExt};
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct FileArgs {
|
||||||
|
path: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct WriteArgs {
|
||||||
|
path: String,
|
||||||
|
content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_request(req: &RpcRequest, root: &Path) -> Result<serde_json::Value, RpcError> {
|
||||||
|
match req.method.as_str() {
|
||||||
|
"initialize" => {
|
||||||
|
let params = req
|
||||||
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| RpcError::invalid_params("Missing params for initialize"))?;
|
||||||
|
|
||||||
|
let init_params: InitializeParams =
|
||||||
|
serde_json::from_value(params.clone()).map_err(|e| {
|
||||||
|
RpcError::invalid_params(format!("Invalid initialize params: {}", e))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Check protocol version compatibility
|
||||||
|
if !is_compatible(&init_params.protocol_version, PROTOCOL_VERSION) {
|
||||||
|
return Err(RpcError::new(
|
||||||
|
ErrorCode::INVALID_REQUEST,
|
||||||
|
format!(
|
||||||
|
"Incompatible protocol version. Client: {}, Server: {}",
|
||||||
|
init_params.protocol_version, PROTOCOL_VERSION
|
||||||
|
),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build initialization result
|
||||||
|
let result = InitializeResult {
|
||||||
|
protocol_version: PROTOCOL_VERSION.to_string(),
|
||||||
|
server_info: ServerInfo {
|
||||||
|
name: "owlen-mcp-server".to_string(),
|
||||||
|
version: env!("CARGO_PKG_VERSION").to_string(),
|
||||||
|
},
|
||||||
|
capabilities: ServerCapabilities {
|
||||||
|
supports_tools: Some(false),
|
||||||
|
supports_resources: Some(true), // Supports read, write, delete
|
||||||
|
supports_streaming: Some(false),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(serde_json::to_value(result).map_err(|e| {
|
||||||
|
RpcError::internal_error(format!("Failed to serialize result: {}", e))
|
||||||
|
})?)
|
||||||
|
}
|
||||||
|
"resources/list" => {
|
||||||
|
let params = req
|
||||||
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| RpcError::invalid_params("Missing params"))?;
|
||||||
|
let args: FileArgs = serde_json::from_value(params.clone())
|
||||||
|
.map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
|
||||||
|
resources_list(&args.path, root).await
|
||||||
|
}
|
||||||
|
"resources/get" => {
|
||||||
|
let params = req
|
||||||
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| RpcError::invalid_params("Missing params"))?;
|
||||||
|
let args: FileArgs = serde_json::from_value(params.clone())
|
||||||
|
.map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
|
||||||
|
resources_get(&args.path, root).await
|
||||||
|
}
|
||||||
|
"resources/write" => {
|
||||||
|
let params = req
|
||||||
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| RpcError::invalid_params("Missing params"))?;
|
||||||
|
let args: WriteArgs = serde_json::from_value(params.clone())
|
||||||
|
.map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
|
||||||
|
resources_write(&args.path, &args.content, root).await
|
||||||
|
}
|
||||||
|
"resources/delete" => {
|
||||||
|
let params = req
|
||||||
|
.params
|
||||||
|
.as_ref()
|
||||||
|
.ok_or_else(|| RpcError::invalid_params("Missing params"))?;
|
||||||
|
let args: FileArgs = serde_json::from_value(params.clone())
|
||||||
|
.map_err(|e| RpcError::invalid_params(format!("Invalid params: {}", e)))?;
|
||||||
|
resources_delete(&args.path, root).await
|
||||||
|
}
|
||||||
|
_ => Err(RpcError::method_not_found(&req.method)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sanitize_path(path: &str, root: &Path) -> Result<PathBuf, RpcError> {
|
||||||
|
let path = Path::new(path);
|
||||||
|
let path = if path.is_absolute() {
|
||||||
|
path.strip_prefix("/")
|
||||||
|
.map_err(|_| RpcError::invalid_params("Invalid path"))?
|
||||||
|
.to_path_buf()
|
||||||
|
} else {
|
||||||
|
path.to_path_buf()
|
||||||
|
};
|
||||||
|
|
||||||
|
let full_path = root.join(path).clean();
|
||||||
|
|
||||||
|
if !full_path.starts_with(root) {
|
||||||
|
return Err(RpcError::path_traversal());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(full_path)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn resources_list(path: &str, root: &Path) -> Result<serde_json::Value, RpcError> {
|
||||||
|
let full_path = sanitize_path(path, root)?;
|
||||||
|
|
||||||
|
let entries = fs::read_dir(full_path).map_err(|e| {
|
||||||
|
RpcError::new(
|
||||||
|
ErrorCode::RESOURCE_NOT_FOUND,
|
||||||
|
format!("Failed to read directory: {}", e),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let mut result = Vec::new();
|
||||||
|
for entry in entries {
|
||||||
|
let entry = entry.map_err(|e| {
|
||||||
|
RpcError::internal_error(format!("Failed to read directory entry: {}", e))
|
||||||
|
})?;
|
||||||
|
result.push(entry.file_name().to_string_lossy().to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(serde_json::json!(result))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn resources_get(path: &str, root: &Path) -> Result<serde_json::Value, RpcError> {
|
||||||
|
let full_path = sanitize_path(path, root)?;
|
||||||
|
|
||||||
|
let content = fs::read_to_string(full_path).map_err(|e| {
|
||||||
|
RpcError::new(
|
||||||
|
ErrorCode::RESOURCE_NOT_FOUND,
|
||||||
|
format!("Failed to read file: {}", e),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok(serde_json::json!(content))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn resources_write(
|
||||||
|
path: &str,
|
||||||
|
content: &str,
|
||||||
|
root: &Path,
|
||||||
|
) -> Result<serde_json::Value, RpcError> {
|
||||||
|
let full_path = sanitize_path(path, root)?;
|
||||||
|
// Ensure parent directory exists
|
||||||
|
if let Some(parent) = full_path.parent() {
|
||||||
|
std::fs::create_dir_all(parent).map_err(|e| {
|
||||||
|
RpcError::internal_error(format!("Failed to create parent directories: {}", e))
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
std::fs::write(full_path, content)
|
||||||
|
.map_err(|e| RpcError::internal_error(format!("Failed to write file: {}", e)))?;
|
||||||
|
Ok(serde_json::json!(null))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn resources_delete(path: &str, root: &Path) -> Result<serde_json::Value, RpcError> {
|
||||||
|
let full_path = sanitize_path(path, root)?;
|
||||||
|
if full_path.is_file() {
|
||||||
|
std::fs::remove_file(full_path)
|
||||||
|
.map_err(|e| RpcError::internal_error(format!("Failed to delete file: {}", e)))?;
|
||||||
|
Ok(serde_json::json!(null))
|
||||||
|
} else {
|
||||||
|
Err(RpcError::new(
|
||||||
|
ErrorCode::RESOURCE_NOT_FOUND,
|
||||||
|
"Path does not refer to a file",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
let root = env::current_dir()?;
|
||||||
|
let mut stdin = io::BufReader::new(io::stdin());
|
||||||
|
let mut stdout = io::stdout();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let mut line = String::new();
|
||||||
|
match stdin.read_line(&mut line).await {
|
||||||
|
Ok(0) => {
|
||||||
|
// EOF
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
Ok(_) => {
|
||||||
|
let req: RpcRequest = match serde_json::from_str(&line) {
|
||||||
|
Ok(req) => req,
|
||||||
|
Err(e) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(
|
||||||
|
RequestId::Number(0),
|
||||||
|
RpcError::parse_error(format!("Parse error: {}", e)),
|
||||||
|
);
|
||||||
|
let resp_str = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(resp_str.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let request_id = req.id.clone();
|
||||||
|
|
||||||
|
match handle_request(&req, &root).await {
|
||||||
|
Ok(result) => {
|
||||||
|
let resp = RpcResponse::new(request_id, result);
|
||||||
|
let resp_str = serde_json::to_string(&resp)?;
|
||||||
|
stdout.write_all(resp_str.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
let err_resp = RpcErrorResponse::new(request_id, error);
|
||||||
|
let resp_str = serde_json::to_string(&err_resp)?;
|
||||||
|
stdout.write_all(resp_str.as_bytes()).await?;
|
||||||
|
stdout.write_all(b"\n").await?;
|
||||||
|
stdout.flush().await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
// Handle read error
|
||||||
|
eprintln!("Error reading from stdin: {}", e);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
This crate provides an implementation of the `owlen-core::Provider` trait for the [Ollama](https://ollama.ai) backend.
|
This crate provides an implementation of the `owlen-core::Provider` trait for the [Ollama](https://ollama.ai) backend.
|
||||||
|
|
||||||
It allows Owlen to communicate with a local Ollama instance, sending requests and receiving responses from locally-run large language models.
|
It allows Owlen to communicate with a local Ollama instance, sending requests and receiving responses from locally-run large language models. You can also target [Ollama Cloud](https://docs.ollama.com/cloud) by pointing the provider at `https://ollama.com` (or `https://api.ollama.com`) and providing an API key through your Owlen configuration (or the `OLLAMA_API_KEY` / `OLLAMA_CLOUD_API_KEY` environment variables). The client automatically adds the required Bearer authorization header when a key is supplied, accepts either host without rewriting, and expands inline environment references like `$OLLAMA_API_KEY` if you prefer not to check the secret into your config file. The generated configuration now includes both `providers.ollama` and `providers.ollama-cloud` entries—switch between them by updating `general.default_provider`.
|
||||||
|
|
||||||
## Configuration
|
## Configuration
|
||||||
|
|
||||||
|
|||||||
@@ -5,13 +5,16 @@ use owlen_core::{
|
|||||||
config::GeneralSettings,
|
config::GeneralSettings,
|
||||||
model::ModelManager,
|
model::ModelManager,
|
||||||
provider::{ChatStream, Provider, ProviderConfig},
|
provider::{ChatStream, Provider, ProviderConfig},
|
||||||
types::{ChatParameters, ChatRequest, ChatResponse, Message, ModelInfo, Role, TokenUsage},
|
types::{
|
||||||
|
ChatParameters, ChatRequest, ChatResponse, Message, ModelInfo, Role, TokenUsage, ToolCall,
|
||||||
|
},
|
||||||
Result,
|
Result,
|
||||||
};
|
};
|
||||||
use reqwest::Client;
|
use reqwest::{header, Client, Url};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
use std::env;
|
||||||
use std::io;
|
use std::io;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
@@ -20,26 +23,195 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
|
|||||||
const DEFAULT_TIMEOUT_SECS: u64 = 120;
|
const DEFAULT_TIMEOUT_SECS: u64 = 120;
|
||||||
const DEFAULT_MODEL_CACHE_TTL_SECS: u64 = 60;
|
const DEFAULT_MODEL_CACHE_TTL_SECS: u64 = 60;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
enum OllamaMode {
|
||||||
|
Local,
|
||||||
|
Cloud,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OllamaMode {
|
||||||
|
fn from_provider_type(provider_type: &str) -> Self {
|
||||||
|
if provider_type.eq_ignore_ascii_case("ollama-cloud") {
|
||||||
|
Self::Cloud
|
||||||
|
} else {
|
||||||
|
Self::Local
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_base_url(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Local => "http://localhost:11434",
|
||||||
|
Self::Cloud => "https://ollama.com",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_scheme(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Local => "http",
|
||||||
|
Self::Cloud => "https",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_ollama_host(host: &str) -> bool {
|
||||||
|
host.eq_ignore_ascii_case("ollama.com")
|
||||||
|
|| host.eq_ignore_ascii_case("www.ollama.com")
|
||||||
|
|| host.eq_ignore_ascii_case("api.ollama.com")
|
||||||
|
|| host.ends_with(".ollama.com")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_base_url(
|
||||||
|
input: Option<&str>,
|
||||||
|
mode_hint: OllamaMode,
|
||||||
|
) -> std::result::Result<String, String> {
|
||||||
|
let mut candidate = input
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(|value| value.to_string())
|
||||||
|
.unwrap_or_else(|| mode_hint.default_base_url().to_string());
|
||||||
|
|
||||||
|
if !candidate.contains("://") {
|
||||||
|
candidate = format!("{}://{}", mode_hint.default_scheme(), candidate);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut url =
|
||||||
|
Url::parse(&candidate).map_err(|err| format!("Invalid base_url '{candidate}': {err}"))?;
|
||||||
|
|
||||||
|
let mut is_cloud = matches!(mode_hint, OllamaMode::Cloud);
|
||||||
|
|
||||||
|
if let Some(host) = url.host_str() {
|
||||||
|
if is_ollama_host(host) {
|
||||||
|
is_cloud = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if is_cloud {
|
||||||
|
if url.scheme() != "https" {
|
||||||
|
url.set_scheme("https")
|
||||||
|
.map_err(|_| "Ollama Cloud requires an https URL".to_string())?;
|
||||||
|
}
|
||||||
|
|
||||||
|
match url.host_str() {
|
||||||
|
Some(host) => {
|
||||||
|
if host.eq_ignore_ascii_case("www.ollama.com") {
|
||||||
|
url.set_host(Some("ollama.com"))
|
||||||
|
.map_err(|_| "Failed to normalize Ollama Cloud host".to_string())?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
return Err("Ollama Cloud base_url must include a hostname".to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove trailing slash and discard query/fragment segments
|
||||||
|
let current_path = url.path().to_string();
|
||||||
|
let trimmed_path = current_path.trim_end_matches('/');
|
||||||
|
if trimmed_path.is_empty() {
|
||||||
|
url.set_path("");
|
||||||
|
} else {
|
||||||
|
url.set_path(trimmed_path);
|
||||||
|
}
|
||||||
|
|
||||||
|
url.set_query(None);
|
||||||
|
url.set_fragment(None);
|
||||||
|
|
||||||
|
Ok(url.to_string().trim_end_matches('/').to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_api_endpoint(base_url: &str, endpoint: &str) -> String {
|
||||||
|
let trimmed_base = base_url.trim_end_matches('/');
|
||||||
|
let trimmed_endpoint = endpoint.trim_start_matches('/');
|
||||||
|
|
||||||
|
if trimmed_base.ends_with("/api") {
|
||||||
|
format!("{trimmed_base}/{trimmed_endpoint}")
|
||||||
|
} else {
|
||||||
|
format!("{trimmed_base}/api/{trimmed_endpoint}")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn env_var_non_empty(name: &str) -> Option<String> {
|
||||||
|
env::var(name)
|
||||||
|
.ok()
|
||||||
|
.map(|value| value.trim().to_string())
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_api_key(configured: Option<String>) -> Option<String> {
|
||||||
|
let raw = configured?.trim().to_string();
|
||||||
|
if raw.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(variable) = raw
|
||||||
|
.strip_prefix("${")
|
||||||
|
.and_then(|value| value.strip_suffix('}'))
|
||||||
|
.or_else(|| raw.strip_prefix('$'))
|
||||||
|
{
|
||||||
|
let var_name = variable.trim();
|
||||||
|
if var_name.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
return env_var_non_empty(var_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
Some(raw)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn debug_requests_enabled() -> bool {
|
||||||
|
std::env::var("OWLEN_DEBUG_OLLAMA")
|
||||||
|
.ok()
|
||||||
|
.map(|value| {
|
||||||
|
matches!(
|
||||||
|
value.trim(),
|
||||||
|
"1" | "true" | "TRUE" | "True" | "yes" | "YES" | "Yes"
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mask_token(token: &str) -> String {
|
||||||
|
if token.len() <= 8 {
|
||||||
|
return "***".to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
let head = &token[..4];
|
||||||
|
let tail = &token[token.len() - 4..];
|
||||||
|
format!("{head}***{tail}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mask_authorization(value: &str) -> String {
|
||||||
|
if let Some(token) = value.strip_prefix("Bearer ") {
|
||||||
|
format!("Bearer {}", mask_token(token))
|
||||||
|
} else {
|
||||||
|
"***".to_string()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Ollama provider implementation with enhanced configuration and caching
|
/// Ollama provider implementation with enhanced configuration and caching
|
||||||
|
#[derive(Debug)]
|
||||||
pub struct OllamaProvider {
|
pub struct OllamaProvider {
|
||||||
client: Client,
|
client: Client,
|
||||||
base_url: String,
|
base_url: String,
|
||||||
|
api_key: Option<String>,
|
||||||
model_manager: ModelManager,
|
model_manager: ModelManager,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Options for configuring the Ollama provider
|
/// Options for configuring the Ollama provider
|
||||||
pub struct OllamaOptions {
|
pub(crate) struct OllamaOptions {
|
||||||
pub base_url: String,
|
base_url: String,
|
||||||
pub request_timeout: Duration,
|
request_timeout: Duration,
|
||||||
pub model_cache_ttl: Duration,
|
model_cache_ttl: Duration,
|
||||||
|
api_key: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl OllamaOptions {
|
impl OllamaOptions {
|
||||||
pub fn new(base_url: impl Into<String>) -> Self {
|
pub(crate) fn new(base_url: impl Into<String>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
base_url: base_url.into(),
|
base_url: base_url.into(),
|
||||||
request_timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
|
request_timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
|
||||||
model_cache_ttl: Duration::from_secs(DEFAULT_MODEL_CACHE_TTL_SECS),
|
model_cache_ttl: Duration::from_secs(DEFAULT_MODEL_CACHE_TTL_SECS),
|
||||||
|
api_key: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -54,6 +226,20 @@ impl OllamaOptions {
|
|||||||
struct OllamaMessage {
|
struct OllamaMessage {
|
||||||
role: String,
|
role: String,
|
||||||
content: String,
|
content: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
tool_calls: Option<Vec<OllamaToolCall>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ollama tool call format
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct OllamaToolCall {
|
||||||
|
function: OllamaToolCallFunction,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct OllamaToolCallFunction {
|
||||||
|
name: String,
|
||||||
|
arguments: serde_json::Value,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Ollama chat request format
|
/// Ollama chat request format
|
||||||
@@ -62,10 +248,27 @@ struct OllamaChatRequest {
|
|||||||
model: String,
|
model: String,
|
||||||
messages: Vec<OllamaMessage>,
|
messages: Vec<OllamaMessage>,
|
||||||
stream: bool,
|
stream: bool,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
tools: Option<Vec<OllamaTool>>,
|
||||||
#[serde(flatten)]
|
#[serde(flatten)]
|
||||||
options: HashMap<String, Value>,
|
options: HashMap<String, Value>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Ollama tool definition
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct OllamaTool {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
tool_type: String,
|
||||||
|
function: OllamaToolFunction,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct OllamaToolFunction {
|
||||||
|
name: String,
|
||||||
|
description: String,
|
||||||
|
parameters: serde_json::Value,
|
||||||
|
}
|
||||||
|
|
||||||
/// Ollama chat response format
|
/// Ollama chat response format
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct OllamaChatResponse {
|
struct OllamaChatResponse {
|
||||||
@@ -107,17 +310,60 @@ struct OllamaModelDetails {
|
|||||||
impl OllamaProvider {
|
impl OllamaProvider {
|
||||||
/// Create a new Ollama provider with sensible defaults
|
/// Create a new Ollama provider with sensible defaults
|
||||||
pub fn new(base_url: impl Into<String>) -> Result<Self> {
|
pub fn new(base_url: impl Into<String>) -> Result<Self> {
|
||||||
Self::with_options(OllamaOptions::new(base_url))
|
let mode = OllamaMode::Local;
|
||||||
|
let supplied = base_url.into();
|
||||||
|
let normalized =
|
||||||
|
normalize_base_url(Some(&supplied), mode).map_err(owlen_core::Error::Config)?;
|
||||||
|
|
||||||
|
Self::with_options(OllamaOptions::new(normalized))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn debug_log_request(&self, label: &str, request: &reqwest::Request, body_json: Option<&str>) {
|
||||||
|
if !debug_requests_enabled() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
eprintln!("--- OWLEN Ollama request ({label}) ---");
|
||||||
|
eprintln!("{} {}", request.method(), request.url());
|
||||||
|
|
||||||
|
match request
|
||||||
|
.headers()
|
||||||
|
.get(header::AUTHORIZATION)
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
{
|
||||||
|
Some(value) => eprintln!("Authorization: {}", mask_authorization(value)),
|
||||||
|
None => eprintln!("Authorization: <none>"),
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(body) = body_json {
|
||||||
|
eprintln!("Body:\n{body}");
|
||||||
|
}
|
||||||
|
|
||||||
|
eprintln!("---------------------------------------");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convert MCP tool descriptors to Ollama tool format
|
||||||
|
fn convert_tools_to_ollama(tools: &[owlen_core::mcp::McpToolDescriptor]) -> Vec<OllamaTool> {
|
||||||
|
tools
|
||||||
|
.iter()
|
||||||
|
.map(|tool| OllamaTool {
|
||||||
|
tool_type: "function".to_string(),
|
||||||
|
function: OllamaToolFunction {
|
||||||
|
name: tool.name.clone(),
|
||||||
|
description: tool.description.clone(),
|
||||||
|
parameters: tool.input_schema.clone(),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create a provider from configuration settings
|
/// Create a provider from configuration settings
|
||||||
pub fn from_config(config: &ProviderConfig, general: Option<&GeneralSettings>) -> Result<Self> {
|
pub fn from_config(config: &ProviderConfig, general: Option<&GeneralSettings>) -> Result<Self> {
|
||||||
let mut options = OllamaOptions::new(
|
let mode = OllamaMode::from_provider_type(&config.provider_type);
|
||||||
config
|
let normalized_base_url = normalize_base_url(config.base_url.as_deref(), mode)
|
||||||
.base_url
|
.map_err(owlen_core::Error::Config)?;
|
||||||
.clone()
|
|
||||||
.unwrap_or_else(|| "http://localhost:11434".to_string()),
|
let mut options = OllamaOptions::new(normalized_base_url);
|
||||||
);
|
|
||||||
|
|
||||||
if let Some(timeout) = config
|
if let Some(timeout) = config
|
||||||
.extra
|
.extra
|
||||||
@@ -135,6 +381,10 @@ impl OllamaProvider {
|
|||||||
options.model_cache_ttl = Duration::from_secs(cache_ttl.max(5));
|
options.model_cache_ttl = Duration::from_secs(cache_ttl.max(5));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
options.api_key = resolve_api_key(config.api_key.clone())
|
||||||
|
.or_else(|| env_var_non_empty("OLLAMA_API_KEY"))
|
||||||
|
.or_else(|| env_var_non_empty("OLLAMA_CLOUD_API_KEY"));
|
||||||
|
|
||||||
if let Some(general) = general {
|
if let Some(general) = general {
|
||||||
options = options.with_general(general);
|
options = options.with_general(general);
|
||||||
}
|
}
|
||||||
@@ -143,16 +393,24 @@ impl OllamaProvider {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Create a provider from explicit options
|
/// Create a provider from explicit options
|
||||||
pub fn with_options(options: OllamaOptions) -> Result<Self> {
|
pub(crate) fn with_options(options: OllamaOptions) -> Result<Self> {
|
||||||
|
let OllamaOptions {
|
||||||
|
base_url,
|
||||||
|
request_timeout,
|
||||||
|
model_cache_ttl,
|
||||||
|
api_key,
|
||||||
|
} = options;
|
||||||
|
|
||||||
let client = Client::builder()
|
let client = Client::builder()
|
||||||
.timeout(options.request_timeout)
|
.timeout(request_timeout)
|
||||||
.build()
|
.build()
|
||||||
.map_err(|e| owlen_core::Error::Config(format!("Failed to build HTTP client: {e}")))?;
|
.map_err(|e| owlen_core::Error::Config(format!("Failed to build HTTP client: {e}")))?;
|
||||||
|
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
client,
|
client,
|
||||||
base_url: options.base_url.trim_end_matches('/').to_string(),
|
base_url: base_url.trim_end_matches('/').to_string(),
|
||||||
model_manager: ModelManager::new(options.model_cache_ttl),
|
api_key,
|
||||||
|
model_manager: ModelManager::new(model_cache_ttl),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -161,14 +419,42 @@ impl OllamaProvider {
|
|||||||
&self.model_manager
|
&self.model_manager
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn api_url(&self, endpoint: &str) -> String {
|
||||||
|
build_api_endpoint(&self.base_url, endpoint)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_auth(&self, request: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
|
||||||
|
if let Some(api_key) = &self.api_key {
|
||||||
|
request.bearer_auth(api_key)
|
||||||
|
} else {
|
||||||
|
request
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn convert_message(message: &Message) -> OllamaMessage {
|
fn convert_message(message: &Message) -> OllamaMessage {
|
||||||
|
let role = match message.role {
|
||||||
|
Role::User => "user".to_string(),
|
||||||
|
Role::Assistant => "assistant".to_string(),
|
||||||
|
Role::System => "system".to_string(),
|
||||||
|
Role::Tool => "tool".to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let tool_calls = message.tool_calls.as_ref().map(|calls| {
|
||||||
|
calls
|
||||||
|
.iter()
|
||||||
|
.map(|tc| OllamaToolCall {
|
||||||
|
function: OllamaToolCallFunction {
|
||||||
|
name: tc.name.clone(),
|
||||||
|
arguments: tc.arguments.clone(),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
OllamaMessage {
|
OllamaMessage {
|
||||||
role: match message.role {
|
role,
|
||||||
Role::User => "user".to_string(),
|
|
||||||
Role::Assistant => "assistant".to_string(),
|
|
||||||
Role::System => "system".to_string(),
|
|
||||||
},
|
|
||||||
content: message.content.clone(),
|
content: message.content.clone(),
|
||||||
|
tool_calls,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -177,10 +463,27 @@ impl OllamaProvider {
|
|||||||
"user" => Role::User,
|
"user" => Role::User,
|
||||||
"assistant" => Role::Assistant,
|
"assistant" => Role::Assistant,
|
||||||
"system" => Role::System,
|
"system" => Role::System,
|
||||||
|
"tool" => Role::Tool,
|
||||||
_ => Role::Assistant,
|
_ => Role::Assistant,
|
||||||
};
|
};
|
||||||
|
|
||||||
Message::new(role, message.content.clone())
|
let mut msg = Message::new(role, message.content.clone());
|
||||||
|
|
||||||
|
// Convert tool calls if present
|
||||||
|
if let Some(ollama_tool_calls) = &message.tool_calls {
|
||||||
|
let tool_calls: Vec<ToolCall> = ollama_tool_calls
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(idx, tc)| ToolCall {
|
||||||
|
id: format!("call_{}", idx),
|
||||||
|
name: tc.function.name.clone(),
|
||||||
|
arguments: tc.function.arguments.clone(),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
msg.tool_calls = Some(tool_calls);
|
||||||
|
}
|
||||||
|
|
||||||
|
msg
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_options(parameters: ChatParameters) -> HashMap<String, Value> {
|
fn build_options(parameters: ChatParameters) -> HashMap<String, Value> {
|
||||||
@@ -202,11 +505,10 @@ impl OllamaProvider {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn fetch_models(&self) -> Result<Vec<ModelInfo>> {
|
async fn fetch_models(&self) -> Result<Vec<ModelInfo>> {
|
||||||
let url = format!("{}/api/tags", self.base_url);
|
let url = self.api_url("tags");
|
||||||
|
|
||||||
let response = self
|
let response = self
|
||||||
.client
|
.apply_auth(self.client.get(&url))
|
||||||
.get(&url)
|
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.map_err(|e| owlen_core::Error::Network(format!("Failed to fetch models: {e}")))?;
|
.map_err(|e| owlen_core::Error::Network(format!("Failed to fetch models: {e}")))?;
|
||||||
@@ -229,21 +531,51 @@ impl OllamaProvider {
|
|||||||
let models = ollama_response
|
let models = ollama_response
|
||||||
.models
|
.models
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|model| ModelInfo {
|
.map(|model| {
|
||||||
id: model.name.clone(),
|
// Check if model supports tool calling based on known models
|
||||||
name: model.name.clone(),
|
let supports_tools = Self::check_tool_support(&model.name);
|
||||||
description: model
|
|
||||||
.details
|
ModelInfo {
|
||||||
.as_ref()
|
id: model.name.clone(),
|
||||||
.and_then(|d| d.family.as_ref().map(|f| format!("Ollama {f} model"))),
|
name: model.name.clone(),
|
||||||
provider: "ollama".to_string(),
|
description: model
|
||||||
context_window: None,
|
.details
|
||||||
capabilities: vec!["chat".to_string()],
|
.as_ref()
|
||||||
|
.and_then(|d| d.family.as_ref().map(|f| format!("Ollama {f} model"))),
|
||||||
|
provider: "ollama".to_string(),
|
||||||
|
context_window: None,
|
||||||
|
capabilities: vec!["chat".to_string()],
|
||||||
|
supports_tools,
|
||||||
|
}
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
Ok(models)
|
Ok(models)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Check if a model supports tool calling based on its name
|
||||||
|
fn check_tool_support(model_name: &str) -> bool {
|
||||||
|
let name_lower = model_name.to_lowercase();
|
||||||
|
|
||||||
|
// Known models with tool calling support
|
||||||
|
let tool_supporting_models = [
|
||||||
|
"qwen",
|
||||||
|
"llama3.1",
|
||||||
|
"llama3.2",
|
||||||
|
"llama3.3",
|
||||||
|
"mistral-nemo",
|
||||||
|
"mistral:7b-instruct",
|
||||||
|
"command-r",
|
||||||
|
"firefunction",
|
||||||
|
"hermes",
|
||||||
|
"nexusraven",
|
||||||
|
"granite-code",
|
||||||
|
];
|
||||||
|
|
||||||
|
tool_supporting_models
|
||||||
|
.iter()
|
||||||
|
.any(|&supported| name_lower.contains(supported))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait::async_trait]
|
#[async_trait::async_trait]
|
||||||
@@ -263,25 +595,52 @@ impl Provider for OllamaProvider {
|
|||||||
model,
|
model,
|
||||||
messages,
|
messages,
|
||||||
parameters,
|
parameters,
|
||||||
|
tools,
|
||||||
} = request;
|
} = request;
|
||||||
|
|
||||||
let messages: Vec<OllamaMessage> = messages.iter().map(Self::convert_message).collect();
|
let messages: Vec<OllamaMessage> = messages.iter().map(Self::convert_message).collect();
|
||||||
|
|
||||||
let options = Self::build_options(parameters);
|
let options = Self::build_options(parameters);
|
||||||
|
|
||||||
|
// Only send the `tools` field if there is at least one tool.
|
||||||
|
// An empty array makes Ollama validate tool support and can cause a
|
||||||
|
// 400 Bad Request for models that do not support tools.
|
||||||
|
// Currently the `tools` field is omitted for compatibility; the variable is retained
|
||||||
|
// for potential future use.
|
||||||
|
let _ollama_tools = tools
|
||||||
|
.as_ref()
|
||||||
|
.filter(|t| !t.is_empty())
|
||||||
|
.map(|t| Self::convert_tools_to_ollama(t));
|
||||||
|
|
||||||
|
// Ollama currently rejects any presence of the `tools` field for models that
|
||||||
|
// do not support function calling. To be safe, we omit the field entirely.
|
||||||
let ollama_request = OllamaChatRequest {
|
let ollama_request = OllamaChatRequest {
|
||||||
model,
|
model,
|
||||||
messages,
|
messages,
|
||||||
stream: false,
|
stream: false,
|
||||||
|
tools: None,
|
||||||
options,
|
options,
|
||||||
};
|
};
|
||||||
|
|
||||||
let url = format!("{}/api/chat", self.base_url);
|
let url = self.api_url("chat");
|
||||||
|
let debug_body = if debug_requests_enabled() {
|
||||||
|
serde_json::to_string_pretty(&ollama_request).ok()
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut request_builder = self.client.post(&url).json(&ollama_request);
|
||||||
|
request_builder = self.apply_auth(request_builder);
|
||||||
|
|
||||||
|
let request = request_builder.build().map_err(|e| {
|
||||||
|
owlen_core::Error::Network(format!("Failed to build chat request: {e}"))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
self.debug_log_request("chat", &request, debug_body.as_deref());
|
||||||
|
|
||||||
let response = self
|
let response = self
|
||||||
.client
|
.client
|
||||||
.post(&url)
|
.execute(request)
|
||||||
.json(&ollama_request)
|
|
||||||
.send()
|
|
||||||
.await
|
.await
|
||||||
.map_err(|e| owlen_core::Error::Network(format!("Chat request failed: {e}")))?;
|
.map_err(|e| owlen_core::Error::Network(format!("Chat request failed: {e}")))?;
|
||||||
|
|
||||||
@@ -339,28 +698,51 @@ impl Provider for OllamaProvider {
|
|||||||
model,
|
model,
|
||||||
messages,
|
messages,
|
||||||
parameters,
|
parameters,
|
||||||
|
tools,
|
||||||
} = request;
|
} = request;
|
||||||
|
|
||||||
let messages: Vec<OllamaMessage> = messages.iter().map(Self::convert_message).collect();
|
let messages: Vec<OllamaMessage> = messages.iter().map(Self::convert_message).collect();
|
||||||
|
|
||||||
let options = Self::build_options(parameters);
|
let options = Self::build_options(parameters);
|
||||||
|
|
||||||
|
// Only include the `tools` field if there is at least one tool.
|
||||||
|
// Sending an empty tools array causes Ollama to reject the request for
|
||||||
|
// models without tool support (400 Bad Request).
|
||||||
|
// Retain tools conversion for possible future extensions, but silence unused warnings.
|
||||||
|
let _ollama_tools = tools
|
||||||
|
.as_ref()
|
||||||
|
.filter(|t| !t.is_empty())
|
||||||
|
.map(|t| Self::convert_tools_to_ollama(t));
|
||||||
|
|
||||||
|
// Omit the `tools` field for compatibility with models lacking tool support.
|
||||||
let ollama_request = OllamaChatRequest {
|
let ollama_request = OllamaChatRequest {
|
||||||
model,
|
model,
|
||||||
messages,
|
messages,
|
||||||
stream: true,
|
stream: true,
|
||||||
|
tools: None,
|
||||||
options,
|
options,
|
||||||
};
|
};
|
||||||
|
|
||||||
let url = format!("{}/api/chat", self.base_url);
|
let url = self.api_url("chat");
|
||||||
|
let debug_body = if debug_requests_enabled() {
|
||||||
|
serde_json::to_string_pretty(&ollama_request).ok()
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
let response = self
|
let mut request_builder = self.client.post(&url).json(&ollama_request);
|
||||||
.client
|
request_builder = self.apply_auth(request_builder);
|
||||||
.post(&url)
|
|
||||||
.json(&ollama_request)
|
let request = request_builder.build().map_err(|e| {
|
||||||
.send()
|
owlen_core::Error::Network(format!("Failed to build streaming request: {e}"))
|
||||||
.await
|
})?;
|
||||||
.map_err(|e| owlen_core::Error::Network(format!("Streaming request failed: {e}")))?;
|
|
||||||
|
self.debug_log_request("chat_stream", &request, debug_body.as_deref());
|
||||||
|
|
||||||
|
let response =
|
||||||
|
self.client.execute(request).await.map_err(|e| {
|
||||||
|
owlen_core::Error::Network(format!("Streaming request failed: {e}"))
|
||||||
|
})?;
|
||||||
|
|
||||||
if !response.status().is_success() {
|
if !response.status().is_success() {
|
||||||
let code = response.status();
|
let code = response.status();
|
||||||
@@ -462,11 +844,10 @@ impl Provider for OllamaProvider {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn health_check(&self) -> Result<()> {
|
async fn health_check(&self) -> Result<()> {
|
||||||
let url = format!("{}/api/version", self.base_url);
|
let url = self.api_url("version");
|
||||||
|
|
||||||
let response = self
|
let response = self
|
||||||
.client
|
.apply_auth(self.client.get(&url))
|
||||||
.get(&url)
|
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.map_err(|e| owlen_core::Error::Network(format!("Health check failed: {e}")))?;
|
.map_err(|e| owlen_core::Error::Network(format!("Health check failed: {e}")))?;
|
||||||
@@ -528,3 +909,86 @@ async fn parse_error_body(response: reqwest::Response) -> String {
|
|||||||
Err(_) => "unknown error".to_string(),
|
Err(_) => "unknown error".to_string(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn normalizes_local_base_url_and_infers_scheme() {
|
||||||
|
let normalized =
|
||||||
|
normalize_base_url(Some("localhost:11434"), OllamaMode::Local).expect("valid URL");
|
||||||
|
assert_eq!(normalized, "http://localhost:11434");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn normalizes_cloud_base_url_and_host() {
|
||||||
|
let normalized =
|
||||||
|
normalize_base_url(Some("https://ollama.com"), OllamaMode::Cloud).expect("valid URL");
|
||||||
|
assert_eq!(normalized, "https://ollama.com");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn infers_scheme_for_cloud_hosts() {
|
||||||
|
let normalized =
|
||||||
|
normalize_base_url(Some("ollama.com"), OllamaMode::Cloud).expect("valid URL");
|
||||||
|
assert_eq!(normalized, "https://ollama.com");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rewrites_www_cloud_host() {
|
||||||
|
let normalized = normalize_base_url(Some("https://www.ollama.com"), OllamaMode::Cloud)
|
||||||
|
.expect("valid URL");
|
||||||
|
assert_eq!(normalized, "https://ollama.com");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn retains_explicit_api_suffix() {
|
||||||
|
let normalized = normalize_base_url(Some("https://api.ollama.com/api"), OllamaMode::Cloud)
|
||||||
|
.expect("valid URL");
|
||||||
|
assert_eq!(normalized, "https://api.ollama.com/api");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn builds_api_endpoint_without_duplicate_segments() {
|
||||||
|
let base = "http://localhost:11434";
|
||||||
|
assert_eq!(
|
||||||
|
build_api_endpoint(base, "chat"),
|
||||||
|
"http://localhost:11434/api/chat"
|
||||||
|
);
|
||||||
|
|
||||||
|
let base_with_api = "http://localhost:11434/api";
|
||||||
|
assert_eq!(
|
||||||
|
build_api_endpoint(base_with_api, "chat"),
|
||||||
|
"http://localhost:11434/api/chat"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn resolve_api_key_prefers_literal_value() {
|
||||||
|
assert_eq!(
|
||||||
|
resolve_api_key(Some("direct-key".into())),
|
||||||
|
Some("direct-key".into())
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn resolve_api_key_expands_braced_env_reference() {
|
||||||
|
std::env::set_var("OWLEN_TEST_KEY_BRACED", "super-secret");
|
||||||
|
assert_eq!(
|
||||||
|
resolve_api_key(Some("${OWLEN_TEST_KEY_BRACED}".into())),
|
||||||
|
Some("super-secret".into())
|
||||||
|
);
|
||||||
|
std::env::remove_var("OWLEN_TEST_KEY_BRACED");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn resolve_api_key_expands_unbraced_env_reference() {
|
||||||
|
std::env::set_var("OWLEN_TEST_KEY_UNBRACED", "another-secret");
|
||||||
|
assert_eq!(
|
||||||
|
resolve_api_key(Some("$OWLEN_TEST_KEY_UNBRACED".into())),
|
||||||
|
Some("another-secret".into())
|
||||||
|
);
|
||||||
|
std::env::remove_var("OWLEN_TEST_KEY_UNBRACED");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -10,6 +10,8 @@ description = "Terminal User Interface for OWLEN LLM client"
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
owlen-core = { path = "../owlen-core" }
|
owlen-core = { path = "../owlen-core" }
|
||||||
|
owlen-ollama = { path = "../owlen-ollama" }
|
||||||
|
# Removed circular dependency on `owlen-cli`. The TUI no longer directly depends on the CLI crate.
|
||||||
|
|
||||||
# TUI framework
|
# TUI framework
|
||||||
ratatui = { workspace = true }
|
ratatui = { workspace = true }
|
||||||
@@ -17,6 +19,7 @@ crossterm = { workspace = true }
|
|||||||
tui-textarea = { workspace = true }
|
tui-textarea = { workspace = true }
|
||||||
textwrap = { workspace = true }
|
textwrap = { workspace = true }
|
||||||
unicode-width = "0.1"
|
unicode-width = "0.1"
|
||||||
|
async-trait = "0.1"
|
||||||
|
|
||||||
# Async runtime
|
# Async runtime
|
||||||
tokio = { workspace = true }
|
tokio = { workspace = true }
|
||||||
@@ -26,6 +29,7 @@ futures-util = { workspace = true }
|
|||||||
# Utilities
|
# Utilities
|
||||||
anyhow = { workspace = true }
|
anyhow = { workspace = true }
|
||||||
uuid = { workspace = true }
|
uuid = { workspace = true }
|
||||||
|
serde_json.workspace = true
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tokio-test = { workspace = true }
|
tokio-test = { workspace = true }
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -14,12 +14,14 @@ pub struct CodeApp {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl CodeApp {
|
impl CodeApp {
|
||||||
pub fn new(mut controller: SessionController) -> (Self, mpsc::UnboundedReceiver<SessionEvent>) {
|
pub async fn new(
|
||||||
|
mut controller: SessionController,
|
||||||
|
) -> Result<(Self, mpsc::UnboundedReceiver<SessionEvent>)> {
|
||||||
controller
|
controller
|
||||||
.conversation_mut()
|
.conversation_mut()
|
||||||
.push_system_message(DEFAULT_SYSTEM_PROMPT.to_string());
|
.push_system_message(DEFAULT_SYSTEM_PROMPT.to_string());
|
||||||
let (inner, rx) = ChatApp::new(controller);
|
let (inner, rx) = ChatApp::new(controller).await?;
|
||||||
(Self { inner }, rx)
|
Ok((Self { inner }, rx))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn handle_event(&mut self, event: Event) -> Result<AppState> {
|
pub async fn handle_event(&mut self, event: Event) -> Result<AppState> {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
pub use owlen_core::config::{
|
pub use owlen_core::config::{
|
||||||
default_config_path, ensure_ollama_config, session_timeout, Config, GeneralSettings,
|
default_config_path, ensure_ollama_config, ensure_provider_config, session_timeout, Config,
|
||||||
InputSettings, StorageSettings, UiSettings, DEFAULT_CONFIG_PATH,
|
GeneralSettings, InputSettings, StorageSettings, UiSettings, DEFAULT_CONFIG_PATH,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Attempt to load configuration from default location
|
/// Attempt to load configuration from default location
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ pub mod chat_app;
|
|||||||
pub mod code_app;
|
pub mod code_app;
|
||||||
pub mod config;
|
pub mod config;
|
||||||
pub mod events;
|
pub mod events;
|
||||||
|
pub mod tui_controller;
|
||||||
pub mod ui;
|
pub mod ui;
|
||||||
|
|
||||||
pub use chat_app::{ChatApp, SessionEvent};
|
pub use chat_app::{ChatApp, SessionEvent};
|
||||||
|
|||||||
44
crates/owlen-tui/src/tui_controller.rs
Normal file
44
crates/owlen-tui/src/tui_controller.rs
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
use async_trait::async_trait;
|
||||||
|
use owlen_core::ui::UiController;
|
||||||
|
use tokio::sync::{mpsc, oneshot};
|
||||||
|
|
||||||
|
/// A request sent from the UiController to the TUI event loop.
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum TuiRequest {
|
||||||
|
Confirm {
|
||||||
|
prompt: String,
|
||||||
|
tx: oneshot::Sender<bool>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
/// An implementation of the UiController trait for the TUI.
|
||||||
|
/// It uses channels to communicate with the main ChatApp event loop.
|
||||||
|
pub struct TuiController {
|
||||||
|
tx: mpsc::UnboundedSender<TuiRequest>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TuiController {
|
||||||
|
pub fn new(tx: mpsc::UnboundedSender<TuiRequest>) -> Self {
|
||||||
|
Self { tx }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl UiController for TuiController {
|
||||||
|
async fn confirm(&self, prompt: &str) -> bool {
|
||||||
|
let (tx, rx) = oneshot::channel();
|
||||||
|
let request = TuiRequest::Confirm {
|
||||||
|
prompt: prompt.to_string(),
|
||||||
|
tx,
|
||||||
|
};
|
||||||
|
|
||||||
|
if self.tx.send(request).is_err() {
|
||||||
|
// Receiver was dropped, so we can't get confirmation.
|
||||||
|
// Default to false for safety.
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for the response from the TUI.
|
||||||
|
rx.await.unwrap_or(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -3,14 +3,17 @@ use ratatui::style::{Color, Modifier, Style};
|
|||||||
use ratatui::text::{Line, Span};
|
use ratatui::text::{Line, Span};
|
||||||
use ratatui::widgets::{Block, Borders, Clear, List, ListItem, ListState, Paragraph, Wrap};
|
use ratatui::widgets::{Block, Borders, Clear, List, ListItem, ListState, Paragraph, Wrap};
|
||||||
use ratatui::Frame;
|
use ratatui::Frame;
|
||||||
|
use serde_json;
|
||||||
use textwrap::{wrap, Options};
|
use textwrap::{wrap, Options};
|
||||||
use tui_textarea::TextArea;
|
use tui_textarea::TextArea;
|
||||||
use unicode_width::UnicodeWidthStr;
|
use unicode_width::UnicodeWidthStr;
|
||||||
|
|
||||||
use crate::chat_app::ChatApp;
|
use crate::chat_app::{ChatApp, ModelSelectorItemKind, HELP_TAB_COUNT};
|
||||||
use owlen_core::types::Role;
|
use owlen_core::types::Role;
|
||||||
use owlen_core::ui::{FocusedPanel, InputMode};
|
use owlen_core::ui::{FocusedPanel, InputMode};
|
||||||
|
|
||||||
|
const PRIVACY_TAB_INDEX: usize = HELP_TAB_COUNT - 1;
|
||||||
|
|
||||||
pub fn render_chat(frame: &mut Frame<'_>, app: &mut ChatApp) {
|
pub fn render_chat(frame: &mut Frame<'_>, app: &mut ChatApp) {
|
||||||
// Update thinking content from last message
|
// Update thinking content from last message
|
||||||
app.update_thinking_from_last_message();
|
app.update_thinking_from_last_message();
|
||||||
@@ -48,6 +51,15 @@ pub fn render_chat(frame: &mut Frame<'_>, app: &mut ChatApp) {
|
|||||||
0
|
0
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Calculate agent actions panel height (similar to thinking)
|
||||||
|
let actions_height = if let Some(actions) = app.agent_actions() {
|
||||||
|
let content_width = available_width.saturating_sub(4);
|
||||||
|
let visual_lines = calculate_wrapped_line_count(actions.lines(), content_width);
|
||||||
|
(visual_lines as u16).min(6) + 2 // +2 for borders, max 6 lines
|
||||||
|
} else {
|
||||||
|
0
|
||||||
|
};
|
||||||
|
|
||||||
let mut constraints = vec![
|
let mut constraints = vec![
|
||||||
Constraint::Length(4), // Header
|
Constraint::Length(4), // Header
|
||||||
Constraint::Min(8), // Messages
|
Constraint::Min(8), // Messages
|
||||||
@@ -56,9 +68,14 @@ pub fn render_chat(frame: &mut Frame<'_>, app: &mut ChatApp) {
|
|||||||
if thinking_height > 0 {
|
if thinking_height > 0 {
|
||||||
constraints.push(Constraint::Length(thinking_height)); // Thinking
|
constraints.push(Constraint::Length(thinking_height)); // Thinking
|
||||||
}
|
}
|
||||||
|
// Insert agent actions panel after thinking (if any)
|
||||||
|
if actions_height > 0 {
|
||||||
|
constraints.push(Constraint::Length(actions_height)); // Agent actions
|
||||||
|
}
|
||||||
|
|
||||||
constraints.push(Constraint::Length(input_height)); // Input
|
constraints.push(Constraint::Length(input_height)); // Input
|
||||||
constraints.push(Constraint::Length(3)); // Status
|
constraints.push(Constraint::Length(5)); // System/Status output (3 lines content + 2 borders)
|
||||||
|
constraints.push(Constraint::Length(3)); // Mode and shortcuts bar
|
||||||
|
|
||||||
let layout = Layout::default()
|
let layout = Layout::default()
|
||||||
.direction(Direction::Vertical)
|
.direction(Direction::Vertical)
|
||||||
@@ -76,20 +93,33 @@ pub fn render_chat(frame: &mut Frame<'_>, app: &mut ChatApp) {
|
|||||||
render_thinking(frame, layout[idx], app);
|
render_thinking(frame, layout[idx], app);
|
||||||
idx += 1;
|
idx += 1;
|
||||||
}
|
}
|
||||||
|
// Render agent actions panel if present
|
||||||
|
if actions_height > 0 {
|
||||||
|
render_agent_actions(frame, layout[idx], app);
|
||||||
|
idx += 1;
|
||||||
|
}
|
||||||
|
|
||||||
render_input(frame, layout[idx], app);
|
render_input(frame, layout[idx], app);
|
||||||
idx += 1;
|
idx += 1;
|
||||||
|
|
||||||
|
render_system_output(frame, layout[idx], app);
|
||||||
|
idx += 1;
|
||||||
|
|
||||||
render_status(frame, layout[idx], app);
|
render_status(frame, layout[idx], app);
|
||||||
|
|
||||||
match app.mode() {
|
// Render consent dialog with highest priority (always on top)
|
||||||
InputMode::ProviderSelection => render_provider_selector(frame, app),
|
if app.has_pending_consent() {
|
||||||
InputMode::ModelSelection => render_model_selector(frame, app),
|
render_consent_dialog(frame, app);
|
||||||
InputMode::Help => render_help(frame, app),
|
} else {
|
||||||
InputMode::SessionBrowser => render_session_browser(frame, app),
|
match app.mode() {
|
||||||
InputMode::ThemeBrowser => render_theme_browser(frame, app),
|
InputMode::ProviderSelection => render_provider_selector(frame, app),
|
||||||
InputMode::Command => render_command_suggestions(frame, app),
|
InputMode::ModelSelection => render_model_selector(frame, app),
|
||||||
_ => {}
|
InputMode::Help => render_help(frame, app),
|
||||||
|
InputMode::SessionBrowser => render_session_browser(frame, app),
|
||||||
|
InputMode::ThemeBrowser => render_theme_browser(frame, app),
|
||||||
|
InputMode::Command => render_command_suggestions(frame, app),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -600,12 +630,16 @@ fn render_messages(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
|||||||
Role::User => ("👤 ", "You: "),
|
Role::User => ("👤 ", "You: "),
|
||||||
Role::Assistant => ("🤖 ", "Assistant: "),
|
Role::Assistant => ("🤖 ", "Assistant: "),
|
||||||
Role::System => ("⚙️ ", "System: "),
|
Role::System => ("⚙️ ", "System: "),
|
||||||
|
Role::Tool => ("🔧 ", "Tool: "),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Extract content without thinking tags for assistant messages
|
// Extract content without thinking tags for assistant messages
|
||||||
let content_to_display = if matches!(role, Role::Assistant) {
|
let content_to_display = if matches!(role, Role::Assistant) {
|
||||||
let (content_without_think, _) = formatter.extract_thinking(&message.content);
|
let (content_without_think, _) = formatter.extract_thinking(&message.content);
|
||||||
content_without_think
|
content_without_think
|
||||||
|
} else if matches!(role, Role::Tool) {
|
||||||
|
// Format tool results nicely
|
||||||
|
format_tool_output(&message.content)
|
||||||
} else {
|
} else {
|
||||||
message.content.clone()
|
message.content.clone()
|
||||||
};
|
};
|
||||||
@@ -658,7 +692,13 @@ fn render_messages(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
|||||||
|
|
||||||
let chunks_len = chunks.len();
|
let chunks_len = chunks.len();
|
||||||
for (i, seg) in chunks.into_iter().enumerate() {
|
for (i, seg) in chunks.into_iter().enumerate() {
|
||||||
let mut spans = vec![Span::raw(format!("{indent}{}", seg))];
|
let style = if matches!(role, Role::Tool) {
|
||||||
|
Style::default().fg(theme.tool_output)
|
||||||
|
} else {
|
||||||
|
Style::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut spans = vec![Span::styled(format!("{indent}{}", seg), style)];
|
||||||
if i == chunks_len - 1 && is_streaming {
|
if i == chunks_len - 1 && is_streaming {
|
||||||
spans.push(Span::styled(" ▌", Style::default().fg(theme.cursor)));
|
spans.push(Span::styled(" ▌", Style::default().fg(theme.cursor)));
|
||||||
}
|
}
|
||||||
@@ -670,7 +710,13 @@ fn render_messages(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
|||||||
let chunks = wrap(&content, content_width as usize);
|
let chunks = wrap(&content, content_width as usize);
|
||||||
let chunks_len = chunks.len();
|
let chunks_len = chunks.len();
|
||||||
for (i, seg) in chunks.into_iter().enumerate() {
|
for (i, seg) in chunks.into_iter().enumerate() {
|
||||||
let mut spans = vec![Span::raw(seg.into_owned())];
|
let style = if matches!(role, Role::Tool) {
|
||||||
|
Style::default().fg(theme.tool_output)
|
||||||
|
} else {
|
||||||
|
Style::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut spans = vec![Span::styled(seg.into_owned(), style)];
|
||||||
if i == chunks_len - 1 && is_streaming {
|
if i == chunks_len - 1 && is_streaming {
|
||||||
spans.push(Span::styled(" ▌", Style::default().fg(theme.cursor)));
|
spans.push(Span::styled(" ▌", Style::default().fg(theme.cursor)));
|
||||||
}
|
}
|
||||||
@@ -870,6 +916,191 @@ fn render_thinking(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Render a panel displaying the latest ReAct agent actions (thought/action/observation).
|
||||||
|
// Color-coded: THOUGHT (blue), ACTION (yellow), OBSERVATION (green)
|
||||||
|
fn render_agent_actions(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
||||||
|
let theme = app.theme().clone();
|
||||||
|
|
||||||
|
if let Some(actions) = app.agent_actions().cloned() {
|
||||||
|
let viewport_height = area.height.saturating_sub(2) as usize; // subtract borders
|
||||||
|
let content_width = area.width.saturating_sub(4);
|
||||||
|
|
||||||
|
// Parse and color-code ReAct components
|
||||||
|
let mut lines: Vec<Line> = Vec::new();
|
||||||
|
|
||||||
|
for line in actions.lines() {
|
||||||
|
let line_trimmed = line.trim();
|
||||||
|
|
||||||
|
// Detect ReAct components and apply color coding
|
||||||
|
if line_trimmed.starts_with("THOUGHT:") {
|
||||||
|
// Blue for THOUGHT
|
||||||
|
let thought_content = line_trimmed.strip_prefix("THOUGHT:").unwrap_or("").trim();
|
||||||
|
let wrapped = wrap(thought_content, content_width as usize);
|
||||||
|
|
||||||
|
// First line with label
|
||||||
|
if let Some(first) = wrapped.first() {
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"THOUGHT: ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Blue)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::styled(first.to_string(), Style::default().fg(Color::Blue)),
|
||||||
|
]));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Continuation lines
|
||||||
|
for chunk in wrapped.iter().skip(1) {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
format!(" {}", chunk),
|
||||||
|
Style::default().fg(Color::Blue),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
} else if line_trimmed.starts_with("ACTION:") {
|
||||||
|
// Yellow for ACTION
|
||||||
|
let action_content = line_trimmed.strip_prefix("ACTION:").unwrap_or("").trim();
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"ACTION: ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Yellow)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::styled(
|
||||||
|
action_content,
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Yellow)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
} else if line_trimmed.starts_with("ACTION_INPUT:") {
|
||||||
|
// Cyan for ACTION_INPUT
|
||||||
|
let input_content = line_trimmed
|
||||||
|
.strip_prefix("ACTION_INPUT:")
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let wrapped = wrap(input_content, content_width as usize);
|
||||||
|
|
||||||
|
if let Some(first) = wrapped.first() {
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"ACTION_INPUT: ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Cyan)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::styled(first.to_string(), Style::default().fg(Color::Cyan)),
|
||||||
|
]));
|
||||||
|
}
|
||||||
|
|
||||||
|
for chunk in wrapped.iter().skip(1) {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
format!(" {}", chunk),
|
||||||
|
Style::default().fg(Color::Cyan),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
} else if line_trimmed.starts_with("OBSERVATION:") {
|
||||||
|
// Green for OBSERVATION
|
||||||
|
let obs_content = line_trimmed
|
||||||
|
.strip_prefix("OBSERVATION:")
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let wrapped = wrap(obs_content, content_width as usize);
|
||||||
|
|
||||||
|
if let Some(first) = wrapped.first() {
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"OBSERVATION: ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Green)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::styled(first.to_string(), Style::default().fg(Color::Green)),
|
||||||
|
]));
|
||||||
|
}
|
||||||
|
|
||||||
|
for chunk in wrapped.iter().skip(1) {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
format!(" {}", chunk),
|
||||||
|
Style::default().fg(Color::Green),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
} else if line_trimmed.starts_with("FINAL_ANSWER:") {
|
||||||
|
// Magenta for FINAL_ANSWER
|
||||||
|
let answer_content = line_trimmed
|
||||||
|
.strip_prefix("FINAL_ANSWER:")
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let wrapped = wrap(answer_content, content_width as usize);
|
||||||
|
|
||||||
|
if let Some(first) = wrapped.first() {
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"FINAL_ANSWER: ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Magenta)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::styled(
|
||||||
|
first.to_string(),
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Magenta)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
}
|
||||||
|
|
||||||
|
for chunk in wrapped.iter().skip(1) {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
format!(" {}", chunk),
|
||||||
|
Style::default().fg(Color::Magenta),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
} else if !line_trimmed.is_empty() {
|
||||||
|
// Regular text
|
||||||
|
let wrapped = wrap(line_trimmed, content_width as usize);
|
||||||
|
for chunk in wrapped {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
chunk.into_owned(),
|
||||||
|
Style::default().fg(theme.text),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Empty line
|
||||||
|
lines.push(Line::from(""));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Highlight border if this panel is focused
|
||||||
|
let border_color = if matches!(app.focused_panel(), FocusedPanel::Thinking) {
|
||||||
|
// Reuse the same focus logic; could add a dedicated enum variant later.
|
||||||
|
theme.focused_panel_border
|
||||||
|
} else {
|
||||||
|
theme.unfocused_panel_border
|
||||||
|
};
|
||||||
|
|
||||||
|
let paragraph = Paragraph::new(lines)
|
||||||
|
.style(Style::default().bg(theme.background))
|
||||||
|
.block(
|
||||||
|
Block::default()
|
||||||
|
.title(Span::styled(
|
||||||
|
" 🤖 Agent Actions ",
|
||||||
|
Style::default()
|
||||||
|
.fg(theme.thinking_panel_title)
|
||||||
|
.add_modifier(Modifier::ITALIC),
|
||||||
|
))
|
||||||
|
.borders(Borders::ALL)
|
||||||
|
.border_style(Style::default().fg(border_color))
|
||||||
|
.style(Style::default().bg(theme.background).fg(theme.text)),
|
||||||
|
)
|
||||||
|
.wrap(Wrap { trim: false });
|
||||||
|
|
||||||
|
frame.render_widget(paragraph, area);
|
||||||
|
_ = viewport_height;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn render_input(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
fn render_input(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
||||||
let theme = app.theme();
|
let theme = app.theme();
|
||||||
let title = match app.mode() {
|
let title = match app.mode() {
|
||||||
@@ -949,6 +1180,47 @@ fn render_input(frame: &mut Frame<'_>, area: Rect, app: &mut ChatApp) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn render_system_output(frame: &mut Frame<'_>, area: Rect, app: &ChatApp) {
|
||||||
|
let theme = app.theme();
|
||||||
|
let system_status = app.system_status();
|
||||||
|
|
||||||
|
// Priority: system_status > error > status > "Ready"
|
||||||
|
let display_message = if !system_status.is_empty() {
|
||||||
|
system_status.to_string()
|
||||||
|
} else if let Some(error) = app.error_message() {
|
||||||
|
format!("Error: {}", error)
|
||||||
|
} else {
|
||||||
|
let status = app.status_message();
|
||||||
|
if status.is_empty() || status == "Ready" {
|
||||||
|
"Ready".to_string()
|
||||||
|
} else {
|
||||||
|
status.to_string()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Create a simple paragraph with wrapping enabled
|
||||||
|
let line = Line::from(Span::styled(
|
||||||
|
display_message,
|
||||||
|
Style::default().fg(theme.info),
|
||||||
|
));
|
||||||
|
|
||||||
|
let paragraph = Paragraph::new(line)
|
||||||
|
.style(Style::default().bg(theme.background))
|
||||||
|
.block(
|
||||||
|
Block::default()
|
||||||
|
.title(Span::styled(
|
||||||
|
" System/Status ",
|
||||||
|
Style::default().fg(theme.info).add_modifier(Modifier::BOLD),
|
||||||
|
))
|
||||||
|
.borders(Borders::ALL)
|
||||||
|
.border_style(Style::default().fg(theme.unfocused_panel_border))
|
||||||
|
.style(Style::default().bg(theme.background).fg(theme.text)),
|
||||||
|
)
|
||||||
|
.wrap(Wrap { trim: false });
|
||||||
|
|
||||||
|
frame.render_widget(paragraph, area);
|
||||||
|
}
|
||||||
|
|
||||||
fn calculate_wrapped_line_count<'a, I>(lines: I, available_width: u16) -> usize
|
fn calculate_wrapped_line_count<'a, I>(lines: I, available_width: u16) -> usize
|
||||||
where
|
where
|
||||||
I: IntoIterator<Item = &'a str>,
|
I: IntoIterator<Item = &'a str>,
|
||||||
@@ -997,39 +1269,39 @@ fn render_status(frame: &mut Frame<'_>, area: Rect, app: &ChatApp) {
|
|||||||
InputMode::ThemeBrowser => (" THEMES", theme.mode_help),
|
InputMode::ThemeBrowser => (" THEMES", theme.mode_help),
|
||||||
};
|
};
|
||||||
|
|
||||||
let status_message = if let Some(error) = app.error_message() {
|
|
||||||
format!("Error: {}", error)
|
|
||||||
} else {
|
|
||||||
app.status_message().to_string()
|
|
||||||
};
|
|
||||||
|
|
||||||
let help_text = "i:Input :m:Model :n:New :c:Clear :h:Help q:Quit";
|
let help_text = "i:Input :m:Model :n:New :c:Clear :h:Help q:Quit";
|
||||||
|
|
||||||
let left_spans = vec![
|
let mut spans = vec![Span::styled(
|
||||||
Span::styled(
|
format!(" {} ", mode_text),
|
||||||
format!(" {} ", mode_text),
|
Style::default()
|
||||||
|
.fg(theme.background)
|
||||||
|
.bg(mode_bg_color)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
)];
|
||||||
|
|
||||||
|
// Add agent status indicator if agent mode is active
|
||||||
|
if app.is_agent_running() {
|
||||||
|
spans.push(Span::styled(
|
||||||
|
" 🤖 AGENT RUNNING ",
|
||||||
Style::default()
|
Style::default()
|
||||||
.fg(theme.background)
|
.fg(Color::Black)
|
||||||
.bg(mode_bg_color)
|
.bg(Color::Yellow)
|
||||||
.add_modifier(Modifier::BOLD),
|
.add_modifier(Modifier::BOLD),
|
||||||
),
|
));
|
||||||
Span::styled(
|
} else if app.is_agent_mode() {
|
||||||
format!(" | {} ", status_message),
|
spans.push(Span::styled(
|
||||||
Style::default().fg(theme.text),
|
" 🤖 AGENT MODE ",
|
||||||
),
|
Style::default()
|
||||||
];
|
.fg(Color::Black)
|
||||||
|
.bg(Color::Cyan)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
let right_spans = vec![
|
spans.push(Span::styled(" ", Style::default().fg(theme.text)));
|
||||||
Span::styled(" Help: ", Style::default().fg(theme.text)),
|
spans.push(Span::styled(help_text, Style::default().fg(theme.info)));
|
||||||
Span::styled(help_text, Style::default().fg(theme.info)),
|
|
||||||
];
|
|
||||||
|
|
||||||
let layout = Layout::default()
|
let paragraph = Paragraph::new(Line::from(spans))
|
||||||
.direction(Direction::Horizontal)
|
|
||||||
.constraints([Constraint::Percentage(50), Constraint::Percentage(50)])
|
|
||||||
.split(area);
|
|
||||||
|
|
||||||
let left_paragraph = Paragraph::new(Line::from(left_spans))
|
|
||||||
.alignment(Alignment::Left)
|
.alignment(Alignment::Left)
|
||||||
.style(Style::default().bg(theme.status_background).fg(theme.text))
|
.style(Style::default().bg(theme.status_background).fg(theme.text))
|
||||||
.block(
|
.block(
|
||||||
@@ -1039,18 +1311,7 @@ fn render_status(frame: &mut Frame<'_>, area: Rect, app: &ChatApp) {
|
|||||||
.style(Style::default().bg(theme.status_background).fg(theme.text)),
|
.style(Style::default().bg(theme.status_background).fg(theme.text)),
|
||||||
);
|
);
|
||||||
|
|
||||||
let right_paragraph = Paragraph::new(Line::from(right_spans))
|
frame.render_widget(paragraph, area);
|
||||||
.alignment(Alignment::Right)
|
|
||||||
.style(Style::default().bg(theme.status_background).fg(theme.text))
|
|
||||||
.block(
|
|
||||||
Block::default()
|
|
||||||
.borders(Borders::ALL)
|
|
||||||
.border_style(Style::default().fg(theme.unfocused_panel_border))
|
|
||||||
.style(Style::default().bg(theme.status_background).fg(theme.text)),
|
|
||||||
);
|
|
||||||
|
|
||||||
frame.render_widget(left_paragraph, layout[0]);
|
|
||||||
frame.render_widget(right_paragraph, layout[1]);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn render_provider_selector(frame: &mut Frame<'_>, app: &ChatApp) {
|
fn render_provider_selector(frame: &mut Frame<'_>, app: &ChatApp) {
|
||||||
@@ -1102,20 +1363,49 @@ fn render_model_selector(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
frame.render_widget(Clear, area);
|
frame.render_widget(Clear, area);
|
||||||
|
|
||||||
let items: Vec<ListItem> = app
|
let items: Vec<ListItem> = app
|
||||||
.models()
|
.model_selector_items()
|
||||||
.iter()
|
.iter()
|
||||||
.map(|model| {
|
.map(|item| match item.kind() {
|
||||||
let label = if model.name.is_empty() {
|
ModelSelectorItemKind::Header { provider, expanded } => {
|
||||||
model.id.clone()
|
let marker = if *expanded { "▼" } else { "▶" };
|
||||||
} else {
|
let label = format!("{} {}", marker, provider);
|
||||||
format!("{} — {}", model.id, model.name)
|
ListItem::new(Span::styled(
|
||||||
};
|
label,
|
||||||
ListItem::new(Span::styled(
|
Style::default()
|
||||||
label,
|
.fg(theme.focused_panel_border)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
ModelSelectorItemKind::Model {
|
||||||
|
provider: _,
|
||||||
|
model_index,
|
||||||
|
} => {
|
||||||
|
if let Some(model) = app.model_info_by_index(*model_index) {
|
||||||
|
let tool_indicator = if model.supports_tools { "🔧 " } else { " " };
|
||||||
|
let label = if model.name.is_empty() {
|
||||||
|
format!(" {}{}", tool_indicator, model.id)
|
||||||
|
} else {
|
||||||
|
format!(" {}{} — {}", tool_indicator, model.id, model.name)
|
||||||
|
};
|
||||||
|
ListItem::new(Span::styled(
|
||||||
|
label,
|
||||||
|
Style::default()
|
||||||
|
.fg(theme.user_message_role)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
ListItem::new(Span::styled(
|
||||||
|
" <model unavailable>",
|
||||||
|
Style::default().fg(theme.error),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ModelSelectorItemKind::Empty { provider } => ListItem::new(Span::styled(
|
||||||
|
format!(" (no models configured for {provider})"),
|
||||||
Style::default()
|
Style::default()
|
||||||
.fg(theme.user_message_role)
|
.fg(theme.unfocused_panel_border)
|
||||||
.add_modifier(Modifier::BOLD),
|
.add_modifier(Modifier::ITALIC),
|
||||||
))
|
)),
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -1123,7 +1413,7 @@ fn render_model_selector(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
.block(
|
.block(
|
||||||
Block::default()
|
Block::default()
|
||||||
.title(Span::styled(
|
.title(Span::styled(
|
||||||
format!("Select Model ({})", app.selected_provider),
|
"Select Model — 🔧 = Tool Support",
|
||||||
Style::default()
|
Style::default()
|
||||||
.fg(theme.focused_panel_border)
|
.fg(theme.focused_panel_border)
|
||||||
.add_modifier(Modifier::BOLD),
|
.add_modifier(Modifier::BOLD),
|
||||||
@@ -1139,10 +1429,232 @@ fn render_model_selector(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
.highlight_symbol("▶ ");
|
.highlight_symbol("▶ ");
|
||||||
|
|
||||||
let mut state = ListState::default();
|
let mut state = ListState::default();
|
||||||
state.select(app.selected_model_index());
|
state.select(app.selected_model_item());
|
||||||
frame.render_stateful_widget(list, area, &mut state);
|
frame.render_stateful_widget(list, area, &mut state);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn render_consent_dialog(frame: &mut Frame<'_>, app: &ChatApp) {
|
||||||
|
let theme = app.theme();
|
||||||
|
|
||||||
|
// Get consent dialog state
|
||||||
|
let consent_state = match app.consent_dialog() {
|
||||||
|
Some(state) => state,
|
||||||
|
None => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Create centered modal area
|
||||||
|
let area = centered_rect(70, 50, frame.area());
|
||||||
|
frame.render_widget(Clear, area);
|
||||||
|
|
||||||
|
// Build consent dialog content
|
||||||
|
let mut lines = vec![
|
||||||
|
Line::from(vec![
|
||||||
|
Span::styled("🔒 ", Style::default().fg(theme.focused_panel_border)),
|
||||||
|
Span::styled(
|
||||||
|
"Consent Required",
|
||||||
|
Style::default()
|
||||||
|
.fg(theme.focused_panel_border)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
Line::from(""),
|
||||||
|
Line::from(vec![
|
||||||
|
Span::styled("Tool: ", Style::default().add_modifier(Modifier::BOLD)),
|
||||||
|
Span::styled(
|
||||||
|
consent_state.tool_name.clone(),
|
||||||
|
Style::default().fg(theme.user_message_role),
|
||||||
|
),
|
||||||
|
]),
|
||||||
|
Line::from(""),
|
||||||
|
];
|
||||||
|
|
||||||
|
// Add data types if any
|
||||||
|
if !consent_state.data_types.is_empty() {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
"Data Access:",
|
||||||
|
Style::default().add_modifier(Modifier::BOLD),
|
||||||
|
)));
|
||||||
|
for data_type in &consent_state.data_types {
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::raw(" • "),
|
||||||
|
Span::styled(data_type, Style::default().fg(theme.text)),
|
||||||
|
]));
|
||||||
|
}
|
||||||
|
lines.push(Line::from(""));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add endpoints if any
|
||||||
|
if !consent_state.endpoints.is_empty() {
|
||||||
|
lines.push(Line::from(Span::styled(
|
||||||
|
"Endpoints:",
|
||||||
|
Style::default().add_modifier(Modifier::BOLD),
|
||||||
|
)));
|
||||||
|
for endpoint in &consent_state.endpoints {
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::raw(" • "),
|
||||||
|
Span::styled(endpoint, Style::default().fg(theme.text)),
|
||||||
|
]));
|
||||||
|
}
|
||||||
|
lines.push(Line::from(""));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add prompt
|
||||||
|
lines.push(Line::from(""));
|
||||||
|
lines.push(Line::from(vec![Span::styled(
|
||||||
|
"Choose consent scope:",
|
||||||
|
Style::default()
|
||||||
|
.fg(theme.focused_panel_border)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
)]));
|
||||||
|
lines.push(Line::from(""));
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"[1] ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Cyan)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::raw("Allow once "),
|
||||||
|
Span::styled(
|
||||||
|
"- Grant only for this operation",
|
||||||
|
Style::default().fg(theme.placeholder),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"[2] ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Green)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::raw("Allow session "),
|
||||||
|
Span::styled(
|
||||||
|
"- Grant for current session",
|
||||||
|
Style::default().fg(theme.placeholder),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"[3] ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::Yellow)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::raw("Allow always "),
|
||||||
|
Span::styled(
|
||||||
|
"- Grant permanently",
|
||||||
|
Style::default().fg(theme.placeholder),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"[4] ",
|
||||||
|
Style::default().fg(Color::Red).add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::raw("Deny "),
|
||||||
|
Span::styled(
|
||||||
|
"- Reject this operation",
|
||||||
|
Style::default().fg(theme.placeholder),
|
||||||
|
),
|
||||||
|
]));
|
||||||
|
lines.push(Line::from(""));
|
||||||
|
lines.push(Line::from(vec![
|
||||||
|
Span::styled(
|
||||||
|
"[Esc] ",
|
||||||
|
Style::default()
|
||||||
|
.fg(Color::DarkGray)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
),
|
||||||
|
Span::raw("Cancel"),
|
||||||
|
]));
|
||||||
|
|
||||||
|
let paragraph = Paragraph::new(lines)
|
||||||
|
.block(
|
||||||
|
Block::default()
|
||||||
|
.title(Span::styled(
|
||||||
|
" Consent Dialog ",
|
||||||
|
Style::default()
|
||||||
|
.fg(theme.focused_panel_border)
|
||||||
|
.add_modifier(Modifier::BOLD),
|
||||||
|
))
|
||||||
|
.borders(Borders::ALL)
|
||||||
|
.border_style(Style::default().fg(theme.focused_panel_border))
|
||||||
|
.style(Style::default().bg(theme.background)),
|
||||||
|
)
|
||||||
|
.alignment(Alignment::Left)
|
||||||
|
.wrap(Wrap { trim: true });
|
||||||
|
|
||||||
|
frame.render_widget(paragraph, area);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn render_privacy_settings(frame: &mut Frame<'_>, area: Rect, app: &ChatApp) {
|
||||||
|
let theme = app.theme();
|
||||||
|
let config = app.config();
|
||||||
|
|
||||||
|
let block = Block::default()
|
||||||
|
.title("Privacy Settings")
|
||||||
|
.borders(Borders::ALL)
|
||||||
|
.border_style(Style::default().fg(theme.unfocused_panel_border))
|
||||||
|
.style(Style::default().bg(theme.background).fg(theme.text));
|
||||||
|
let inner = block.inner(area);
|
||||||
|
frame.render_widget(block, area);
|
||||||
|
|
||||||
|
let remote_search_enabled =
|
||||||
|
config.privacy.enable_remote_search && config.tools.web_search.enabled;
|
||||||
|
let code_exec_enabled = config.tools.code_exec.enabled;
|
||||||
|
let history_days = config.privacy.retain_history_days;
|
||||||
|
let cache_results = config.privacy.cache_web_results;
|
||||||
|
let consent_required = config.privacy.require_consent_per_session;
|
||||||
|
let encryption_enabled = config.privacy.encrypt_local_data;
|
||||||
|
|
||||||
|
let status_line = |label: &str, enabled: bool| {
|
||||||
|
let status_text = if enabled { "Enabled" } else { "Disabled" };
|
||||||
|
let status_style = if enabled {
|
||||||
|
Style::default().fg(theme.selection_fg)
|
||||||
|
} else {
|
||||||
|
Style::default().fg(theme.error)
|
||||||
|
};
|
||||||
|
Line::from(vec![
|
||||||
|
Span::raw(format!(" {label}: ")),
|
||||||
|
Span::styled(status_text, status_style),
|
||||||
|
])
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut lines = Vec::new();
|
||||||
|
lines.push(Line::from(vec![Span::styled(
|
||||||
|
"Privacy Configuration",
|
||||||
|
Style::default().fg(theme.info).add_modifier(Modifier::BOLD),
|
||||||
|
)]));
|
||||||
|
lines.push(Line::raw(""));
|
||||||
|
lines.push(Line::from("Network Access:"));
|
||||||
|
lines.push(status_line("Web Search", remote_search_enabled));
|
||||||
|
lines.push(status_line("Code Execution", code_exec_enabled));
|
||||||
|
lines.push(Line::raw(""));
|
||||||
|
lines.push(Line::from("Data Retention:"));
|
||||||
|
lines.push(Line::from(format!(
|
||||||
|
" History retention: {} day(s)",
|
||||||
|
history_days
|
||||||
|
)));
|
||||||
|
lines.push(Line::from(format!(
|
||||||
|
" Cache web results: {}",
|
||||||
|
if cache_results { "Yes" } else { "No" }
|
||||||
|
)));
|
||||||
|
lines.push(Line::raw(""));
|
||||||
|
lines.push(Line::from("Safeguards:"));
|
||||||
|
lines.push(status_line("Consent required", consent_required));
|
||||||
|
lines.push(status_line("Encrypted storage", encryption_enabled));
|
||||||
|
lines.push(Line::raw(""));
|
||||||
|
lines.push(Line::from("Commands:"));
|
||||||
|
lines.push(Line::from(" :privacy-enable <tool> - Enable tool"));
|
||||||
|
lines.push(Line::from(" :privacy-disable <tool> - Disable tool"));
|
||||||
|
lines.push(Line::from(" :privacy-clear - Clear all data"));
|
||||||
|
|
||||||
|
let paragraph = Paragraph::new(lines)
|
||||||
|
.wrap(Wrap { trim: true })
|
||||||
|
.style(Style::default().bg(theme.background).fg(theme.text));
|
||||||
|
frame.render_widget(paragraph, inner);
|
||||||
|
}
|
||||||
|
|
||||||
fn render_help(frame: &mut Frame<'_>, app: &ChatApp) {
|
fn render_help(frame: &mut Frame<'_>, app: &ChatApp) {
|
||||||
let theme = app.theme();
|
let theme = app.theme();
|
||||||
let area = centered_rect(75, 70, frame.area());
|
let area = centered_rect(75, 70, frame.area());
|
||||||
@@ -1156,6 +1668,7 @@ fn render_help(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
"Commands",
|
"Commands",
|
||||||
"Sessions",
|
"Sessions",
|
||||||
"Browsers",
|
"Browsers",
|
||||||
|
"Privacy",
|
||||||
];
|
];
|
||||||
|
|
||||||
// Build tab line
|
// Build tab line
|
||||||
@@ -1429,6 +1942,7 @@ fn render_help(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
Line::from(" g / Home → jump to top"),
|
Line::from(" g / Home → jump to top"),
|
||||||
Line::from(" G / End → jump to bottom"),
|
Line::from(" G / End → jump to bottom"),
|
||||||
],
|
],
|
||||||
|
6 => vec![],
|
||||||
|
|
||||||
_ => vec![],
|
_ => vec![],
|
||||||
};
|
};
|
||||||
@@ -1454,14 +1968,18 @@ fn render_help(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
frame.render_widget(tabs_para, layout[0]);
|
frame.render_widget(tabs_para, layout[0]);
|
||||||
|
|
||||||
// Render content
|
// Render content
|
||||||
let content_block = Block::default()
|
if tab_index == PRIVACY_TAB_INDEX {
|
||||||
.borders(Borders::LEFT | Borders::RIGHT)
|
render_privacy_settings(frame, layout[1], app);
|
||||||
.border_style(Style::default().fg(theme.unfocused_panel_border))
|
} else {
|
||||||
.style(Style::default().bg(theme.background).fg(theme.text));
|
let content_block = Block::default()
|
||||||
let content_para = Paragraph::new(help_text)
|
.borders(Borders::LEFT | Borders::RIGHT)
|
||||||
.style(Style::default().bg(theme.background).fg(theme.text))
|
.border_style(Style::default().fg(theme.unfocused_panel_border))
|
||||||
.block(content_block);
|
.style(Style::default().bg(theme.background).fg(theme.text));
|
||||||
frame.render_widget(content_para, layout[1]);
|
let content_para = Paragraph::new(help_text)
|
||||||
|
.style(Style::default().bg(theme.background).fg(theme.text))
|
||||||
|
.block(content_block);
|
||||||
|
frame.render_widget(content_para, layout[1]);
|
||||||
|
}
|
||||||
|
|
||||||
// Render navigation hint
|
// Render navigation hint
|
||||||
let nav_hint = Line::from(vec![
|
let nav_hint = Line::from(vec![
|
||||||
@@ -1474,7 +1992,7 @@ fn render_help(frame: &mut Frame<'_>, app: &ChatApp) {
|
|||||||
),
|
),
|
||||||
Span::raw(":Switch "),
|
Span::raw(":Switch "),
|
||||||
Span::styled(
|
Span::styled(
|
||||||
"1-6",
|
format!("1-{}", HELP_TAB_COUNT),
|
||||||
Style::default()
|
Style::default()
|
||||||
.fg(theme.focused_panel_border)
|
.fg(theme.focused_panel_border)
|
||||||
.add_modifier(Modifier::BOLD),
|
.add_modifier(Modifier::BOLD),
|
||||||
@@ -1845,6 +2363,108 @@ fn role_color(role: &Role, theme: &owlen_core::theme::Theme) -> Style {
|
|||||||
match role {
|
match role {
|
||||||
Role::User => Style::default().fg(theme.user_message_role),
|
Role::User => Style::default().fg(theme.user_message_role),
|
||||||
Role::Assistant => Style::default().fg(theme.assistant_message_role),
|
Role::Assistant => Style::default().fg(theme.assistant_message_role),
|
||||||
Role::System => Style::default().fg(theme.info),
|
Role::System => Style::default().fg(theme.unfocused_panel_border),
|
||||||
|
Role::Tool => Style::default().fg(theme.info),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Format tool output JSON into a nice human-readable format
|
||||||
|
fn format_tool_output(content: &str) -> String {
|
||||||
|
// Try to parse as JSON
|
||||||
|
if let Ok(json) = serde_json::from_str::<serde_json::Value>(content) {
|
||||||
|
let mut output = String::new();
|
||||||
|
let mut content_found = false;
|
||||||
|
|
||||||
|
// Extract query if present
|
||||||
|
if let Some(query) = json.get("query").and_then(|v| v.as_str()) {
|
||||||
|
output.push_str(&format!("Query: \"{}\"\n\n", query));
|
||||||
|
content_found = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Extract results array
|
||||||
|
if let Some(results) = json.get("results").and_then(|v| v.as_array()) {
|
||||||
|
content_found = true;
|
||||||
|
if results.is_empty() {
|
||||||
|
output.push_str("No results found");
|
||||||
|
return output;
|
||||||
|
}
|
||||||
|
|
||||||
|
for (i, result) in results.iter().enumerate() {
|
||||||
|
// Title
|
||||||
|
if let Some(title) = result.get("title").and_then(|v| v.as_str()) {
|
||||||
|
// Strip HTML tags from title
|
||||||
|
let clean_title = title.replace("<b>", "").replace("</b>", "");
|
||||||
|
output.push_str(&format!("{}. {}\n", i + 1, clean_title));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Source and date (if available)
|
||||||
|
let mut meta = Vec::new();
|
||||||
|
if let Some(source) = result.get("source").and_then(|v| v.as_str()) {
|
||||||
|
meta.push(format!("📰 {}", source));
|
||||||
|
}
|
||||||
|
if let Some(date) = result.get("date").and_then(|v| v.as_str()) {
|
||||||
|
// Simplify date format
|
||||||
|
if let Some(simple_date) = date.split('T').next() {
|
||||||
|
meta.push(format!("📅 {}", simple_date));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !meta.is_empty() {
|
||||||
|
output.push_str(&format!(" {}\n", meta.join(" • ")));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Snippet (truncated if too long)
|
||||||
|
if let Some(snippet) = result.get("snippet").and_then(|v| v.as_str()) {
|
||||||
|
if !snippet.is_empty() {
|
||||||
|
// Strip HTML tags
|
||||||
|
let clean_snippet = snippet
|
||||||
|
.replace("<b>", "")
|
||||||
|
.replace("</b>", "")
|
||||||
|
.replace("'", "'")
|
||||||
|
.replace(""", "\"");
|
||||||
|
|
||||||
|
// Truncate if too long
|
||||||
|
let truncated = if clean_snippet.len() > 200 {
|
||||||
|
format!("{}...", &clean_snippet[..197])
|
||||||
|
} else {
|
||||||
|
clean_snippet
|
||||||
|
};
|
||||||
|
output.push_str(&format!(" {}\n", truncated));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// URL (shortened if too long)
|
||||||
|
if let Some(url) = result.get("url").and_then(|v| v.as_str()) {
|
||||||
|
let display_url = if url.len() > 80 {
|
||||||
|
format!("{}...", &url[..77])
|
||||||
|
} else {
|
||||||
|
url.to_string()
|
||||||
|
};
|
||||||
|
output.push_str(&format!(" 🔗 {}\n", display_url));
|
||||||
|
}
|
||||||
|
|
||||||
|
output.push('\n');
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add total count
|
||||||
|
if let Some(total) = json.get("total_found").and_then(|v| v.as_u64()) {
|
||||||
|
output.push_str(&format!("Found {} result(s)", total));
|
||||||
|
}
|
||||||
|
} else if let Some(result) = json.get("result").and_then(|v| v.as_str()) {
|
||||||
|
content_found = true;
|
||||||
|
output.push_str(result);
|
||||||
|
} else if let Some(error) = json.get("error").and_then(|v| v.as_str()) {
|
||||||
|
content_found = true;
|
||||||
|
// Handle error results
|
||||||
|
output.push_str(&format!("❌ Error: {}", error));
|
||||||
|
}
|
||||||
|
|
||||||
|
if content_found {
|
||||||
|
output
|
||||||
|
} else {
|
||||||
|
content.to_string()
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// If not JSON, return as-is
|
||||||
|
content.to_string()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -96,17 +96,22 @@ These settings control the behavior of the text input area.
|
|||||||
|
|
||||||
## Provider Settings (`[providers]`)
|
## Provider Settings (`[providers]`)
|
||||||
|
|
||||||
This section contains a table for each provider you want to configure. The key is the provider name (e.g., `ollama`).
|
This section contains a table for each provider you want to configure. Owlen ships with two entries pre-populated: `ollama` for a local daemon and `ollama-cloud` for the hosted API. You can switch between them by changing `general.default_provider`.
|
||||||
|
|
||||||
```toml
|
```toml
|
||||||
[providers.ollama]
|
[providers.ollama]
|
||||||
provider_type = "ollama"
|
provider_type = "ollama"
|
||||||
base_url = "http://localhost:11434"
|
base_url = "http://localhost:11434"
|
||||||
# api_key = "..."
|
# api_key = "..."
|
||||||
|
|
||||||
|
[providers.ollama-cloud]
|
||||||
|
provider_type = "ollama-cloud"
|
||||||
|
base_url = "https://ollama.com"
|
||||||
|
# api_key = "${OLLAMA_API_KEY}"
|
||||||
```
|
```
|
||||||
|
|
||||||
- `provider_type` (string, required)
|
- `provider_type` (string, required)
|
||||||
The type of the provider. Currently, only `"ollama"` is built-in.
|
The type of the provider. The built-in options are `"ollama"` (local daemon) and `"ollama-cloud"` (hosted service).
|
||||||
|
|
||||||
- `base_url` (string, optional)
|
- `base_url` (string, optional)
|
||||||
The base URL of the provider's API.
|
The base URL of the provider's API.
|
||||||
@@ -116,3 +121,16 @@ base_url = "http://localhost:11434"
|
|||||||
|
|
||||||
- `extra` (table, optional)
|
- `extra` (table, optional)
|
||||||
Any additional, provider-specific parameters can be added here.
|
Any additional, provider-specific parameters can be added here.
|
||||||
|
|
||||||
|
### Using Ollama Cloud
|
||||||
|
|
||||||
|
To talk to [Ollama Cloud](https://docs.ollama.com/cloud), point the base URL at the hosted endpoint and supply your API key:
|
||||||
|
|
||||||
|
```toml
|
||||||
|
[providers.ollama-cloud]
|
||||||
|
provider_type = "ollama-cloud"
|
||||||
|
base_url = "https://ollama.com"
|
||||||
|
api_key = "${OLLAMA_API_KEY}"
|
||||||
|
```
|
||||||
|
|
||||||
|
Requests target the same `/api/chat` endpoint documented by Ollama and automatically include the API key using a `Bearer` authorization header. If you prefer not to store the key in the config file, you can leave `api_key` unset and provide it via the `OLLAMA_API_KEY` (or `OLLAMA_CLOUD_API_KEY`) environment variable instead. You can also reference an environment variable inline (for example `api_key = "$OLLAMA_API_KEY"` or `api_key = "${OLLAMA_API_KEY}"`), which Owlen expands when the configuration is loaded. The base URL is normalised automatically—Owlen enforces HTTPS, trims trailing slashes, and accepts both `https://ollama.com` and `https://api.ollama.com` without rewriting the host.
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "lightmagenta"
|
|||||||
unfocused_panel_border = "#5f1487"
|
unfocused_panel_border = "#5f1487"
|
||||||
user_message_role = "lightblue"
|
user_message_role = "lightblue"
|
||||||
assistant_message_role = "yellow"
|
assistant_message_role = "yellow"
|
||||||
|
tool_output = "gray"
|
||||||
thinking_panel_title = "lightmagenta"
|
thinking_panel_title = "lightmagenta"
|
||||||
command_bar_background = "black"
|
command_bar_background = "black"
|
||||||
status_background = "black"
|
status_background = "black"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#4a90e2"
|
|||||||
unfocused_panel_border = "#dddddd"
|
unfocused_panel_border = "#dddddd"
|
||||||
user_message_role = "#0055a4"
|
user_message_role = "#0055a4"
|
||||||
assistant_message_role = "#8e44ad"
|
assistant_message_role = "#8e44ad"
|
||||||
|
tool_output = "gray"
|
||||||
thinking_panel_title = "#8e44ad"
|
thinking_panel_title = "#8e44ad"
|
||||||
command_bar_background = "white"
|
command_bar_background = "white"
|
||||||
status_background = "white"
|
status_background = "white"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#ff79c6"
|
|||||||
unfocused_panel_border = "#44475a"
|
unfocused_panel_border = "#44475a"
|
||||||
user_message_role = "#8be9fd"
|
user_message_role = "#8be9fd"
|
||||||
assistant_message_role = "#ff79c6"
|
assistant_message_role = "#ff79c6"
|
||||||
|
tool_output = "#6272a4"
|
||||||
thinking_panel_title = "#bd93f9"
|
thinking_panel_title = "#bd93f9"
|
||||||
command_bar_background = "#44475a"
|
command_bar_background = "#44475a"
|
||||||
status_background = "#44475a"
|
status_background = "#44475a"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#fe8019"
|
|||||||
unfocused_panel_border = "#7c6f64"
|
unfocused_panel_border = "#7c6f64"
|
||||||
user_message_role = "#b8bb26"
|
user_message_role = "#b8bb26"
|
||||||
assistant_message_role = "#83a598"
|
assistant_message_role = "#83a598"
|
||||||
|
tool_output = "#928374"
|
||||||
thinking_panel_title = "#d3869b"
|
thinking_panel_title = "#d3869b"
|
||||||
command_bar_background = "#3c3836"
|
command_bar_background = "#3c3836"
|
||||||
status_background = "#3c3836"
|
status_background = "#3c3836"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#80cbc4"
|
|||||||
unfocused_panel_border = "#546e7a"
|
unfocused_panel_border = "#546e7a"
|
||||||
user_message_role = "#82aaff"
|
user_message_role = "#82aaff"
|
||||||
assistant_message_role = "#c792ea"
|
assistant_message_role = "#c792ea"
|
||||||
|
tool_output = "#546e7a"
|
||||||
thinking_panel_title = "#ffcb6b"
|
thinking_panel_title = "#ffcb6b"
|
||||||
command_bar_background = "#212b30"
|
command_bar_background = "#212b30"
|
||||||
status_background = "#212b30"
|
status_background = "#212b30"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#009688"
|
|||||||
unfocused_panel_border = "#b0bec5"
|
unfocused_panel_border = "#b0bec5"
|
||||||
user_message_role = "#448aff"
|
user_message_role = "#448aff"
|
||||||
assistant_message_role = "#7c4dff"
|
assistant_message_role = "#7c4dff"
|
||||||
|
tool_output = "#90a4ae"
|
||||||
thinking_panel_title = "#f57c00"
|
thinking_panel_title = "#f57c00"
|
||||||
command_bar_background = "#ffffff"
|
command_bar_background = "#ffffff"
|
||||||
status_background = "#ffffff"
|
status_background = "#ffffff"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#58a6ff"
|
|||||||
unfocused_panel_border = "#30363d"
|
unfocused_panel_border = "#30363d"
|
||||||
user_message_role = "#79c0ff"
|
user_message_role = "#79c0ff"
|
||||||
assistant_message_role = "#89ddff"
|
assistant_message_role = "#89ddff"
|
||||||
|
tool_output = "#546e7a"
|
||||||
thinking_panel_title = "#9ece6a"
|
thinking_panel_title = "#9ece6a"
|
||||||
command_bar_background = "#161b22"
|
command_bar_background = "#161b22"
|
||||||
status_background = "#161b22"
|
status_background = "#161b22"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#f92672"
|
|||||||
unfocused_panel_border = "#75715e"
|
unfocused_panel_border = "#75715e"
|
||||||
user_message_role = "#66d9ef"
|
user_message_role = "#66d9ef"
|
||||||
assistant_message_role = "#ae81ff"
|
assistant_message_role = "#ae81ff"
|
||||||
|
tool_output = "#75715e"
|
||||||
thinking_panel_title = "#e6db74"
|
thinking_panel_title = "#e6db74"
|
||||||
command_bar_background = "#272822"
|
command_bar_background = "#272822"
|
||||||
status_background = "#272822"
|
status_background = "#272822"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#eb6f92"
|
|||||||
unfocused_panel_border = "#26233a"
|
unfocused_panel_border = "#26233a"
|
||||||
user_message_role = "#31748f"
|
user_message_role = "#31748f"
|
||||||
assistant_message_role = "#9ccfd8"
|
assistant_message_role = "#9ccfd8"
|
||||||
|
tool_output = "#6e6a86"
|
||||||
thinking_panel_title = "#c4a7e7"
|
thinking_panel_title = "#c4a7e7"
|
||||||
command_bar_background = "#26233a"
|
command_bar_background = "#26233a"
|
||||||
status_background = "#26233a"
|
status_background = "#26233a"
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ focused_panel_border = "#268bd2"
|
|||||||
unfocused_panel_border = "#073642"
|
unfocused_panel_border = "#073642"
|
||||||
user_message_role = "#2aa198"
|
user_message_role = "#2aa198"
|
||||||
assistant_message_role = "#cb4b16"
|
assistant_message_role = "#cb4b16"
|
||||||
|
tool_output = "#657b83"
|
||||||
thinking_panel_title = "#6c71c4"
|
thinking_panel_title = "#6c71c4"
|
||||||
command_bar_background = "#073642"
|
command_bar_background = "#073642"
|
||||||
status_background = "#073642"
|
status_background = "#073642"
|
||||||
|
|||||||
Reference in New Issue
Block a user