vrr/quill-code
2
0
Fork 0
mirror of https://github.com/taggart-comet/quill-code.git synced 2026-10-19 09:44:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

save progress

Browse source
This commit is contained in:
taggart_comet 2026-02-13 19:56:43 +02:00
parent 95b32ed25c
commit a86ad6e41c
55 changed files with 2763 additions and 818 deletions
Download patch file Download diff file Expand all files Collapse all files

531
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

6
Cargo.toml
View file

@ -66,6 +66,12 @@ pulldown-cmark = "0.10"
arboard = "3.6"
base64 = "0.22"
image = { version = "0.25", default-features = false, features = ["png", "jpeg"] }
oauth2 = "4.4"
tiny_http = "0.12"
open = "5.0"
sha2 = "0.10"
url = "2.5"
rand = "0.8"
[dev-dependencies]

2
cats/cat.py
View file

@ -14,7 +14,7 @@ class Cat:
def pet(self) -> str:
if self.mood == "happy":
return f"{self.name} purrs loudly."
return f"{self.name} tolerates the petting."
return f"{self.name} tolerates the petting with a calm blink."
def begs(self) -> str:
self.mood = "needy"

42
src/domain/mod.rs
View file

@ -5,20 +5,20 @@ mod project;
pub mod prompting;
pub mod session;
pub mod startup;
pub mod todo;
pub mod tools;
mod user_settings;
pub mod workflow;
pub mod todo;
pub use agent_mode::AgentModeType;
pub use project::Project;
pub use session::service::SessionService;
pub use session::{Session, SessionRequest};
pub use startup::StartupService;
pub use user_settings::UserSettings;
pub use workflow::{CancellationToken, Chain};
#[allow(unused_imports)]
pub use todo::{TodoList, TodoItem};
pub use todo::{TodoItem, TodoList};
pub use user_settings::{AuthMethod, UserSettings};
pub use workflow::{CancellationToken, Chain};
/// Model type enum matching the inference engine types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -27,6 +27,13 @@ pub enum ModelType {
OpenAI,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelAuthType {
Local,
ApiKey,
OAuth,
}
impl ModelType {
pub fn as_str(&self) -> &'static str {
match self {
@ -42,4 +49,29 @@ impl ModelType {
_ => None,
}
}
}
}
impl ModelAuthType {
pub fn as_str(&self) -> &'static str {
match self {
ModelAuthType::Local => "local",
ModelAuthType::ApiKey => "api_key",
ModelAuthType::OAuth => "oauth",
}
}
pub fn from_str(value: &str) -> Self {
match value {
"local" => ModelAuthType::Local,
"oauth" => ModelAuthType::OAuth,
_ => ModelAuthType::ApiKey,
}
}
pub fn from_auth_method(method: &AuthMethod) -> Self {
match method {
AuthMethod::ApiKey => ModelAuthType::ApiKey,
AuthMethod::OAuth => ModelAuthType::OAuth,
}
}
}

468
src/domain/permissions/checker.rs
View file

@ -1,8 +1,10 @@
use super::store::{PermissionStore, StoreError};
use super::types::{PermissionConfig, PermissionDecision, PermissionRequest, PermissionScope};
use super::types::{
PermissionConfig, PermissionRequest, PermissionScope, SystemPermissionDecision,
UserPermissionDecision,
};
use crate::domain::session::Request;
use crate::domain::tools::Tool;
use crate::utils::paths::is_within_root;
use crate::utils::AskError;
use std::path::PathBuf;
use std::sync::Arc;
@ -23,7 +25,10 @@ pub struct PermissionChecker {
}
pub trait PermissionPrompter: Send + Sync {
fn ask_permission(&self, request: &PermissionRequest) -> Result<PermissionDecision, AskError>;
fn ask_permission(
&self,
request: &PermissionRequest,
) -> Result<UserPermissionDecision, AskError>;
}
impl PermissionChecker {
@ -46,60 +51,68 @@ impl PermissionChecker {
request: &dyn Request,
project_id: Option<i32>,
) -> Result<bool, CheckerError> {
// Auto-allow read-only operations within project root
if tool.is_read_only() {
let all_within_project = tool.get_affected_paths(request)
.iter()
.all(|p| p.starts_with(&request.project_root().to_path_buf()));
if all_within_project {
return Ok(true);
}
}
let permission_request = PermissionRequest::new(
tool.name().to_string(),
tool.get_command(request),
tool.get_affected_paths(request),
super::types::PermissionScope::Project,
PermissionScope::Project,
project_id,
tool.is_read_only(),
request.project_root().to_path_buf(),
);
let decision = self.resolve_permission(&permission_request)?;
if decision == PermissionDecision::Ask
&& self.is_allowed_by_default(tool, request, &permission_request)
{
return Ok(true);
}
match decision {
PermissionDecision::AllowOnce => Ok(true),
PermissionDecision::Ask => self.prompt_and_store(&permission_request),
PermissionDecision::AllowAllReadsInSession
| PermissionDecision::AllowAllWritesInSession
| PermissionDecision::AllowCommandForProject => {
// These should have been resolved already, but handle them safely
Ok(true)
// Dangerous commands always require a prompt
if let Some(ref cmd) = permission_request.command {
if self.is_dangerous_command(cmd) {
return self.prompt_and_store(&permission_request);
}
}
// Restricted paths always require a prompt
if permission_request
.paths
.iter()
.any(|p| self.is_restricted_path(p))
{
return self.prompt_and_store(&permission_request);
}
let decision = self.resolve_permission(&permission_request)?;
match decision {
SystemPermissionDecision::Allow => Ok(true),
SystemPermissionDecision::Ask => self.prompt_and_store(&permission_request),
}
}
/// Store a permission decision
fn store_permission_decision(
&self,
request: &PermissionRequest,
decision: PermissionDecision,
user_decision: UserPermissionDecision,
) -> Result<(), CheckerError> {
// Handle command-specific permissions (project-scoped, persistent)
if decision == PermissionDecision::AllowCommandForProject {
let permission = super::types::Permission::new(
request.tool_name.clone(),
request.command.clone(), // Store the specific command
None, // No resource pattern for commands
decision,
PermissionScope::Project, // Project scope = persistent
request.project_id,
);
self.store.create_permission(permission)?;
if user_decision == UserPermissionDecision::AllowOnce {
return Ok(());
}
// Handle session-wide permissions
if !matches!(decision, PermissionDecision::AllowAllReadsInSession | PermissionDecision::AllowAllWritesInSession) {
return Ok(());
}
// Get project root pattern for session-wide permissions
let resource_pattern = if let Some(first_path) = request.paths.first() {
// Minimal behavior: for project-scoped patch_files AlwaysAllow, store without
// resource pattern so it applies to any path within the project for this tool.
let resource_pattern = if request.tool_name == "patch_files"
&& user_decision == UserPermissionDecision::AlwaysAllow
&& request.scope == PermissionScope::Project
{
None
} else if let Some(first_path) = request.paths.first() {
// Try to get the project root by going up the directory tree
let mut path = first_path.clone();
while let Some(parent) = path.parent() {
@ -116,10 +129,10 @@ impl PermissionChecker {
let permission = super::types::Permission::new(
request.tool_name.clone(),
None, // No command pattern for session-wide
None, // No command pattern for session-wide
resource_pattern,
decision,
PermissionScope::Session,
user_decision,
PermissionScope::Project,
request.project_id,
);
@ -144,123 +157,38 @@ impl PermissionChecker {
.any(|restricted| path_str.starts_with(restricted) || path_str.contains(restricted))
}
fn check_default_rules(
&self,
request: &PermissionRequest,
) -> Result<PermissionDecision, CheckerError> {
// Check for dangerous commands
if let Some(command) = &request.command {
if self.is_dangerous_command(command) {
return Ok(PermissionDecision::Ask);
}
}
// Check for restricted paths
for path in &request.paths {
if self.is_restricted_path(path) {
return Ok(PermissionDecision::Ask);
}
}
// Use default decision
Ok(self.config.default_decision.clone())
}
fn is_allowed_by_default(
&self,
tool: &dyn Tool,
request: &dyn Request,
permission_request: &PermissionRequest,
) -> bool {
// Allow safe tools that don't require permission
if tool.skip_permission_check() {
return true;
}
// Allow read-only tools if paths are within project root
if !tool.is_read_only() {
return false;
}
let project_root = request.project_root();
if permission_request.paths.is_empty() {
return false;
}
permission_request.paths.iter().all(|path| {
let normalized = if path.is_absolute() {
path.to_path_buf()
} else {
project_root.join(path)
};
is_within_root(&normalized, project_root)
})
}
fn resolve_permission(
&self,
request: &PermissionRequest,
) -> Result<PermissionDecision, CheckerError> {
) -> Result<SystemPermissionDecision, CheckerError> {
let project_id: i32 = request.project_id.unwrap_or(0);
if project_id == 0 {
return Ok(PermissionDecision::Ask);
return Ok(SystemPermissionDecision::Ask);
}
// Check for command-specific permissions first (project-scoped, persistent)
if let Some(command) = &request.command {
if let Ok(Some(perm)) = self.store.find_permission(
&request.tool_name,
project_id,
command,
"",
) {
if perm.decision == PermissionDecision::AllowCommandForProject
&& perm.scope == PermissionScope::Project {
return Ok(PermissionDecision::AllowOnce); // Grant this specific operation
}
}
let command_str = request.command.as_ref().map_or("", |s| s.as_str());
let path_str = request
.paths
.first()
.map(|p| p.to_string_lossy().to_string())
.unwrap_or_default();
let perm =
self.store
.find_permission(&request.tool_name, project_id, command_str, &path_str);
if let Ok(Some(p)) = perm {
return Ok(p.system_decision());
}
// Check for session-wide permissions
let session_permissions = self.store.find_session_permissions(project_id)?;
for perm in session_permissions {
// CRITICAL SECURITY CHECK: Only grant session-wide permissions if ALL paths are within project root
let all_paths_in_project = request.paths.iter().all(|path| {
let normalized = if path.is_absolute() {
path.to_path_buf()
} else {
request.project_root.join(path)
};
is_within_root(&normalized, &request.project_root)
});
if !all_paths_in_project {
// Paths outside project - don't use session permission, ask explicitly
continue;
}
if request.is_read_only && perm.is_session_wide_all_reads() {
return Ok(PermissionDecision::AllowOnce); // Grant this specific operation
}
if !request.is_read_only && perm.is_session_wide_all_writes() {
return Ok(PermissionDecision::AllowOnce); // Grant this specific operation
}
}
// Apply default security rules
self.check_default_rules(request)
Ok(self.config.default_decision.clone())
}
fn prompt_and_store(&self, request: &PermissionRequest) -> Result<bool, CheckerError> {
match self.prompter.ask_permission(request) {
Ok(decision @ PermissionDecision::AllowAllReadsInSession)
| Ok(decision @ PermissionDecision::AllowAllWritesInSession)
| Ok(decision @ PermissionDecision::AllowCommandForProject) => {
Ok(decision) => {
self.store_permission_decision(request, decision)?;
Ok(true)
}
Ok(PermissionDecision::AllowOnce) => Ok(true),
Ok(PermissionDecision::Ask) => Ok(false),
Err(AskError::IoError) => {
Err(CheckerError::Failed("Permission prompt failed".to_string()))
}
@ -272,7 +200,6 @@ impl PermissionChecker {
mod tests {
use super::*;
use crate::domain::permissions::types::Permission;
use crate::domain::permissions::types::{PermissionConfig, PermissionDecision};
use crate::domain::permissions::{PermissionRequest, PermissionScope};
use crate::domain::session::{Request, SessionRequest};
use crate::domain::tools::{Error, Tool, ToolResult};
@ -307,24 +234,6 @@ mod tests {
Ok(permission)
}
fn find_session_permissions(&self, project_id: i32) -> Result<Vec<Permission>, StoreError> {
let permissions = self.permissions.lock().unwrap();
let session_perms: Vec<Permission> = permissions
.iter()
.filter(|p| {
p.project_id == Some(project_id)
&& p.scope == PermissionScope::Session
&& matches!(
p.decision,
PermissionDecision::AllowAllReadsInSession
| PermissionDecision::AllowAllWritesInSession
)
})
.cloned()
.collect();
Ok(session_perms)
}
fn find_permission(
&self,
tool: &str,
@ -400,14 +309,14 @@ mod tests {
struct TestPrompter {
calls: Arc<AtomicUsize>,
decision: PermissionDecision,
decision: UserPermissionDecision,
}
impl PermissionPrompter for TestPrompter {
fn ask_permission(
&self,
_request: &PermissionRequest,
) -> Result<PermissionDecision, AskError> {
) -> Result<UserPermissionDecision, AskError> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(self.decision.clone())
}
@ -598,12 +507,12 @@ mod tests {
let calls = Arc::new(AtomicUsize::new(0));
let prompter = Arc::new(TestPrompter {
calls: Arc::clone(&calls),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
});
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::Ask,
default_decision: SystemPermissionDecision::Ask,
..PermissionConfig::default()
},
prompter,
@ -631,12 +540,12 @@ mod tests {
let calls = Arc::new(AtomicUsize::new(0));
let prompter = Arc::new(TestPrompter {
calls: Arc::clone(&calls),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
});
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::Ask,
default_decision: SystemPermissionDecision::Ask,
..PermissionConfig::default()
},
prompter,
@ -663,12 +572,12 @@ mod tests {
let calls = Arc::new(AtomicUsize::new(0));
let prompter = Arc::new(TestPrompter {
calls: Arc::clone(&calls),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
});
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::Ask,
default_decision: SystemPermissionDecision::Ask,
..PermissionConfig::default()
},
prompter,
@ -693,12 +602,12 @@ mod tests {
let calls = Arc::new(AtomicUsize::new(0));
let prompter = Arc::new(TestPrompter {
calls: Arc::clone(&calls),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
});
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::AllowOnce,
default_decision: SystemPermissionDecision::Allow,
..PermissionConfig::default()
},
prompter,
@ -725,12 +634,12 @@ mod tests {
let calls = Arc::new(AtomicUsize::new(0));
let prompter = Arc::new(TestPrompter {
calls: Arc::clone(&calls),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
});
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::AllowOnce,
default_decision: SystemPermissionDecision::Allow,
..PermissionConfig::default()
},
prompter,
@ -748,7 +657,7 @@ mod tests {
}
#[test]
fn resolve_permission_returns_ask_for_read_only() {
fn resolve_permission_returns_ask_when_no_stored_permission() {
let abs = std::fs::canonicalize(".").expect("failed to canonicalize path");
let store = Arc::new(TestStore::with_permissions(vec![]));
@ -757,7 +666,7 @@ mod tests {
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
let request = PermissionRequest::new(
@ -766,13 +675,13 @@ mod tests {
vec![abs.clone()],
PermissionScope::Project,
Some(1),
true, // is_read_only
true,
abs.parent().unwrap_or(&abs).to_path_buf(),
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, PermissionDecision::Ask);
assert_eq!(decision, SystemPermissionDecision::Ask);
}
#[test]
@ -781,19 +690,21 @@ mod tests {
let root = temp.path().to_path_buf();
let file_path = root.join("sample.txt");
// Command permission with AllowOnce → system_decision = Ask
let command_permission = Permission::new(
"command_tool".to_string(),
Some("rm -rf /".to_string()),
None,
PermissionDecision::Ask,
UserPermissionDecision::AllowOnce,
PermissionScope::Project,
Some(1),
);
// Path permission with AlwaysAllow → system_decision = Allow
let path_permission = Permission::new(
"command_tool".to_string(),
None,
Some(file_path.to_string_lossy().to_string()),
PermissionDecision::AllowOnce,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
@ -806,7 +717,7 @@ mod tests {
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
let request = PermissionRequest::new(
@ -815,13 +726,14 @@ mod tests {
vec![file_path.clone()],
PermissionScope::Project,
Some(1),
false, // is_read_only
false,
root,
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, PermissionDecision::Ask);
// Command match has higher specificity, and AllowOnce → Ask
assert_eq!(decision, SystemPermissionDecision::Ask);
}
#[test]
@ -834,7 +746,7 @@ mod tests {
"command_tool".to_string(),
None,
Some(file_path.to_string_lossy().to_string()),
PermissionDecision::AllowOnce,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
@ -844,7 +756,7 @@ mod tests {
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
let request = PermissionRequest::new(
@ -853,13 +765,14 @@ mod tests {
vec![file_path.clone()],
PermissionScope::Project,
Some(1),
false, // is_read_only
false,
root,
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, PermissionDecision::Ask);
// Path permission with AlwaysAllow → Allow
assert_eq!(decision, SystemPermissionDecision::Allow);
}
#[test]
@ -868,7 +781,7 @@ mod tests {
"command_tool".to_string(),
None,
None,
PermissionDecision::Ask,
UserPermissionDecision::AllowOnce,
PermissionScope::Project,
Some(1),
);
@ -878,7 +791,7 @@ mod tests {
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
let request = PermissionRequest::new(
@ -887,13 +800,14 @@ mod tests {
vec![],
PermissionScope::Project,
Some(1),
false, // is_read_only
false,
PathBuf::from("/tmp"),
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, PermissionDecision::Ask);
// AllowOnce → system_decision = Ask
assert_eq!(decision, SystemPermissionDecision::Ask);
}
#[test]
@ -902,12 +816,12 @@ mod tests {
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::AllowOnce,
default_decision: SystemPermissionDecision::Allow,
..PermissionConfig::default()
},
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
let request = PermissionRequest::new(
@ -916,147 +830,181 @@ mod tests {
vec![],
PermissionScope::Project,
Some(1),
false, // is_read_only
false,
PathBuf::from("/tmp"),
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, PermissionDecision::AllowOnce);
assert_eq!(decision, SystemPermissionDecision::Allow);
}
#[test]
fn session_read_permission_only_works_within_project() {
fn stored_permission_does_not_apply_to_different_resource() {
let project_dir = tempfile::tempdir().unwrap();
let external_dir = tempfile::tempdir().unwrap();
let project_root = project_dir.path().to_path_buf();
let external_file = external_dir.path().join("external.txt");
std::fs::write(&external_file, "data").unwrap();
// Create a session-wide read permission
let session_perm = Permission::new(
// Create a permission with a specific resource pattern
let perm = Permission::new(
"read_only".to_string(),
None,
Some(format!("{}/**", project_root.to_string_lossy())),
PermissionDecision::AllowAllReadsInSession,
PermissionScope::Session,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
let store = Arc::new(TestStore::with_permissions(vec![session_perm]));
let store = Arc::new(TestStore::with_permissions(vec![perm]));
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
// Try to read a file OUTSIDE the project
// Request for a file that doesn't match the resource pattern
let request = PermissionRequest::new(
"read_only".to_string(),
None,
vec![external_file],
PermissionScope::Project,
Some(1),
true, // is_read_only
true,
project_root,
);
let decision = checker.resolve_permission(&request).unwrap();
// Should NOT grant access via session permission - should Ask
assert_eq!(decision, PermissionDecision::Ask);
// Resource pattern doesn't match → falls back to default (Ask)
assert_eq!(decision, SystemPermissionDecision::Ask);
}
#[test]
fn session_write_permission_only_works_within_project() {
fn stored_permission_does_not_apply_to_different_project() {
let project_dir = tempfile::tempdir().unwrap();
let external_dir = tempfile::tempdir().unwrap();
let project_root = project_dir.path().to_path_buf();
let external_file = external_dir.path().join("external.txt");
std::fs::write(&external_file, "data").unwrap();
// Create a session-wide write permission
let session_perm = Permission::new(
// Create a permission for project 1
let perm = Permission::new(
"write_tool".to_string(),
None,
Some(format!("{}/**", project_root.to_string_lossy())),
PermissionDecision::AllowAllWritesInSession,
PermissionScope::Session,
None,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
let store = Arc::new(TestStore::with_permissions(vec![session_perm]));
let store = Arc::new(TestStore::with_permissions(vec![perm]));
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
// Try to write a file OUTSIDE the project
// Request for project 2 - should not match
let request = PermissionRequest::new(
"write_tool".to_string(),
None,
vec![external_file],
vec![project_root.join("file.txt")],
PermissionScope::Project,
Some(1),
false, // is_read_only
Some(2),
false,
project_root,
);
let decision = checker.resolve_permission(&request).unwrap();
// Should NOT grant access via session permission - should Ask
assert_eq!(decision, PermissionDecision::Ask);
assert_eq!(decision, SystemPermissionDecision::Ask);
}
#[test]
fn session_read_permission_works_within_project() {
fn stored_write_permission_grants_access() {
let project_dir = tempfile::tempdir().unwrap();
let project_root = project_dir.path().to_path_buf();
// Create an AlwaysAllow permission for write_tool (no resource pattern = matches any)
let perm = Permission::new(
"write_tool".to_string(),
None,
None,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
let store = Arc::new(TestStore::with_permissions(vec![perm]));
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: UserPermissionDecision::AlwaysAllow,
}),
);
let request = PermissionRequest::new(
"write_tool".to_string(),
None,
vec![project_root.join("internal.txt")],
PermissionScope::Project,
Some(1),
false,
project_root,
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, SystemPermissionDecision::Allow);
}
#[test]
fn stored_read_permission_grants_access() {
let project_dir = tempfile::tempdir().unwrap();
let project_root = project_dir.path().to_path_buf();
let internal_file = project_root.join("internal.txt");
std::fs::write(&internal_file, "data").unwrap();
// Create a session-wide read permission
let session_perm = Permission::new(
// Create an AlwaysAllow permission for read_only (no resource pattern = matches any)
let perm = Permission::new(
"read_only".to_string(),
None,
Some(format!("{}/**", project_root.to_string_lossy())),
PermissionDecision::AllowAllReadsInSession,
PermissionScope::Session,
None,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
let store = Arc::new(TestStore::with_permissions(vec![session_perm]));
let store = Arc::new(TestStore::with_permissions(vec![perm]));
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
// Try to read a file INSIDE the project
let request = PermissionRequest::new(
"read_only".to_string(),
None,
vec![internal_file],
PermissionScope::Project,
Some(1),
true, // is_read_only
true,
project_root,
);
let decision = checker.resolve_permission(&request).unwrap();
// SHOULD grant access via session permission
assert_eq!(decision, PermissionDecision::AllowOnce);
assert_eq!(decision, SystemPermissionDecision::Allow);
}
#[test]
@ -1064,12 +1012,12 @@ mod tests {
let project_dir = tempfile::tempdir().unwrap();
let project_root = project_dir.path().to_path_buf();
// Create a project-scoped command permission
// Create a project-scoped command permission with AlwaysAllow
let command_perm = Permission::new(
"shell_exec".to_string(),
Some("npm test".to_string()),
None,
PermissionDecision::AllowCommandForProject,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
@ -1080,7 +1028,7 @@ mod tests {
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AlwaysAllow,
}),
);
@ -1091,12 +1039,12 @@ mod tests {
vec![],
PermissionScope::Project,
Some(1),
false, // is_read_only
false,
project_root.clone(),
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, PermissionDecision::AllowOnce); // Allowed via stored permission
assert_eq!(decision, SystemPermissionDecision::Allow);
// Request for a DIFFERENT command should ask
let different_request = PermissionRequest::new(
@ -1110,6 +1058,48 @@ mod tests {
);
let decision = checker.resolve_permission(&different_request).unwrap();
assert_eq!(decision, PermissionDecision::Ask); // Different command, must ask
assert_eq!(decision, SystemPermissionDecision::Ask);
}
#[test]
fn patch_files_permission_within_project_root_applies_to_other_project_files() {
let project_dir = tempfile::tempdir().unwrap();
let project_root = project_dir.path().to_path_buf();
// User previously created an AlwaysAllow permission for patch_files
// targeting one file inside this project.
let existing_perm = Permission::new(
"patch_files".to_string(),
None,
None,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
let store = Arc::new(TestStore::with_permissions(vec![existing_perm]));
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig::default(),
Arc::new(TestPrompter {
calls: Arc::new(AtomicUsize::new(0)),
decision: UserPermissionDecision::AlwaysAllow,
}),
);
// New request patches a DIFFERENT file, but still within the same project root.
// Expected behavior (desired): this should be allowed by the existing permission.
let request = PermissionRequest::new(
"patch_files".to_string(),
None,
vec![project_root.join("src/b.rs")],
PermissionScope::Project,
Some(1),
false,
project_root,
);
let decision = checker.resolve_permission(&request).unwrap();
assert_eq!(decision, SystemPermissionDecision::Allow);
}
}

152
src/domain/permissions/store.rs
View file

@ -1,4 +1,4 @@
use super::types::{Permission, PermissionDecision, PermissionScope};
use super::types::{Permission, PermissionScope, UserPermissionDecision};
use crate::infrastructure::db::DbPool;
use rusqlite::params;
use thiserror::Error;
@ -18,7 +18,6 @@ pub trait PermissionStore {
command_pattern: &str,
resource_pattern: &str,
) -> Result<Option<Permission>, StoreError>;
fn find_session_permissions(&self, project_id: i32) -> Result<Vec<Permission>, StoreError>;
}
pub struct SqlitePermissionStore {
@ -31,18 +30,6 @@ impl SqlitePermissionStore {
}
fn row_to_permission(&self, row: &rusqlite::Row) -> Result<Permission, rusqlite::Error> {
let decision_str: String = row.get("decision")?;
let decision = match decision_str.as_str() {
"allow_all_reads_session" => PermissionDecision::AllowAllReadsInSession,
"allow_all_writes_session" => PermissionDecision::AllowAllWritesInSession,
"allow_command_for_project" => PermissionDecision::AllowCommandForProject,
"ask" => PermissionDecision::Ask,
"once" => PermissionDecision::AllowOnce,
// Handle legacy values gracefully
"allow" | "deny" => PermissionDecision::Ask,
_ => PermissionDecision::Ask, // Default to ask on error
};
let scope_str: String = row.get("scope")?;
let scope = match scope_str.as_str() {
"session" => PermissionScope::Session,
@ -61,7 +48,7 @@ impl SqlitePermissionStore {
tool_name: row.get("tool_name")?,
command_pattern: row.get("command_pattern")?,
resource_pattern: row.get("resource_pattern")?,
decision,
user_decision: UserPermissionDecision::AlwaysAllow,
scope,
project_id: row.get("project_id")?,
created_at,
@ -71,12 +58,10 @@ impl SqlitePermissionStore {
impl PermissionStore for SqlitePermissionStore {
fn create_permission(&self, permission: Permission) -> Result<Permission, StoreError> {
let decision_str = match permission.decision {
PermissionDecision::AllowAllReadsInSession => "allow_all_reads_session",
PermissionDecision::AllowAllWritesInSession => "allow_all_writes_session",
PermissionDecision::AllowCommandForProject => "allow_command_for_project",
PermissionDecision::Ask => "ask",
PermissionDecision::AllowOnce => "once",
let decision_str = match permission.user_decision {
UserPermissionDecision::AlwaysAllow => "always_allow",
UserPermissionDecision::AllowOnce => "allow_once",
UserPermissionDecision::Deny => "deny",
};
let scope_str = match permission.scope {
@ -159,10 +144,8 @@ impl PermissionStore for SqlitePermissionStore {
})?;
// Build params list dynamically to match the query
let mut param_values: Vec<Box<dyn rusqlite::ToSql>> = vec![
Box::new(project_id),
Box::new(tool.to_string()),
];
let mut param_values: Vec<Box<dyn rusqlite::ToSql>> =
vec![Box::new(project_id), Box::new(tool.to_string())];
if !command_pattern.is_empty() {
param_values.push(Box::new(command_pattern.to_string()));
@ -172,10 +155,8 @@ impl PermissionStore for SqlitePermissionStore {
param_values.push(Box::new(resource_pattern.to_string()));
}
let params_refs: Vec<&dyn rusqlite::ToSql> = param_values
.iter()
.map(|p| p.as_ref())
.collect();
let params_refs: Vec<&dyn rusqlite::ToSql> =
param_values.iter().map(|p| p.as_ref()).collect();
let result = conn.query_row(&query, params_refs.as_slice(), |row| {
self.row_to_permission(row)
@ -187,32 +168,6 @@ impl PermissionStore for SqlitePermissionStore {
Err(e) => Err(StoreError::Database(e)),
}
}
fn find_session_permissions(&self, project_id: i32) -> Result<Vec<Permission>, StoreError> {
let conn = self.conn.get().map_err(|e| {
StoreError::Database(rusqlite::Error::ToSqlConversionFailure(Box::new(
std::io::Error::new(
std::io::ErrorKind::Other,
format!("Failed to get connection: {}", e),
),
)))
})?;
let mut stmt = conn.prepare(
"SELECT id, tool_name, command_pattern, resource_pattern, decision, scope,
project_id, created_at
FROM permissions
WHERE project_id = ?1
AND scope = 'session'
AND decision IN ('allow_all_reads_session', 'allow_all_writes_session')",
)?;
let permissions = stmt
.query_map(params![project_id], |row| self.row_to_permission(row))?
.collect::<Result<Vec<_>, _>>()?;
Ok(permissions)
}
}
#[cfg(test)]
@ -238,7 +193,8 @@ mod tests {
created_at TEXT NOT NULL
)",
[],
).unwrap();
)
.unwrap();
pool
}
@ -253,7 +209,7 @@ mod tests {
"test_tool".to_string(),
None,
None,
PermissionDecision::AllowOnce,
UserPermissionDecision::AllowOnce,
PermissionScope::Project,
Some(1),
);
@ -262,7 +218,10 @@ mod tests {
// Should find the permission when searching with empty strings
let result = store.find_permission("test_tool", 1, "", "").unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap().decision, PermissionDecision::AllowOnce);
assert_eq!(
result.unwrap().user_decision,
UserPermissionDecision::AlwaysAllow
);
}
#[test]
@ -275,16 +234,21 @@ mod tests {
"test_tool".to_string(),
Some("echo test".to_string()),
None,
PermissionDecision::Ask,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
store.create_permission(permission).unwrap();
// Should find the permission when searching with command and empty resource
let result = store.find_permission("test_tool", 1, "echo test", "").unwrap();
let result = store
.find_permission("test_tool", 1, "echo test", "")
.unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap().decision, PermissionDecision::Ask);
assert_eq!(
result.unwrap().user_decision,
UserPermissionDecision::AlwaysAllow
);
}
#[test]
@ -297,16 +261,21 @@ mod tests {
"test_tool".to_string(),
None,
Some("/path/to/file".to_string()),
PermissionDecision::AllowOnce,
UserPermissionDecision::AllowOnce,
PermissionScope::Project,
Some(1),
);
store.create_permission(permission).unwrap();
// Should find the permission when searching with empty command and resource
let result = store.find_permission("test_tool", 1, "", "/path/to/file").unwrap();
let result = store
.find_permission("test_tool", 1, "", "/path/to/file")
.unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap().decision, PermissionDecision::AllowOnce);
assert_eq!(
result.unwrap().user_decision,
UserPermissionDecision::AlwaysAllow
);
}
#[test]
@ -319,16 +288,21 @@ mod tests {
"test_tool".to_string(),
Some("rm -rf".to_string()),
Some("/etc/passwd".to_string()),
PermissionDecision::Ask,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
store.create_permission(permission).unwrap();
// Should find the permission when searching with both patterns
let result = store.find_permission("test_tool", 1, "rm -rf", "/etc/passwd").unwrap();
let result = store
.find_permission("test_tool", 1, "rm -rf", "/etc/passwd")
.unwrap();
assert!(result.is_some());
assert_eq!(result.unwrap().decision, PermissionDecision::Ask);
assert_eq!(
result.unwrap().user_decision,
UserPermissionDecision::AlwaysAllow
);
}
#[test]
@ -341,7 +315,7 @@ mod tests {
"test_tool".to_string(),
Some("echo test".to_string()),
None,
PermissionDecision::AllowOnce,
UserPermissionDecision::AllowOnce,
PermissionScope::Project,
Some(1),
);
@ -353,47 +327,23 @@ mod tests {
}
#[test]
fn test_find_session_permissions() {
fn test_find_permission_different_project_no_match() {
let pool = setup_test_db();
let store = SqlitePermissionStore::new(pool.clone());
// Create session-wide read permission
let read_perm = Permission::new(
// Create permission for project 1
let permission = Permission::new(
"read_tool".to_string(),
None,
Some("/project/**".to_string()),
PermissionDecision::AllowAllReadsInSession,
PermissionScope::Session,
Some(1),
);
store.create_permission(read_perm).unwrap();
// Create session-wide write permission
let write_perm = Permission::new(
"write_tool".to_string(),
None,
Some("/project/**".to_string()),
PermissionDecision::AllowAllWritesInSession,
PermissionScope::Session,
Some(1),
);
store.create_permission(write_perm).unwrap();
// Create non-session permission (should not be returned)
let other_perm = Permission::new(
"other_tool".to_string(),
None,
None,
PermissionDecision::AllowOnce,
UserPermissionDecision::AlwaysAllow,
PermissionScope::Project,
Some(1),
);
store.create_permission(other_perm).unwrap();
store.create_permission(permission).unwrap();
// Find session permissions
let result = store.find_session_permissions(1).unwrap();
assert_eq!(result.len(), 2);
assert!(result.iter().any(|p| matches!(p.decision, PermissionDecision::AllowAllReadsInSession)));
assert!(result.iter().any(|p| matches!(p.decision, PermissionDecision::AllowAllWritesInSession)));
// Should NOT find when searching for project 2
let result = store.find_permission("read_tool", 2, "", "").unwrap();
assert!(result.is_none());
}
}
}

111
src/domain/permissions/types.rs
View file

@ -3,12 +3,16 @@ use serde::{Deserialize, Serialize};
use std::path::PathBuf;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum PermissionDecision {
Ask, // Prompt user again
AllowOnce, // One-time approval
AllowAllReadsInSession, // Session-wide read access
AllowAllWritesInSession, // Session-wide write access
AllowCommandForProject, // Allow specific command for this project (persistent)
pub enum SystemPermissionDecision {
Ask,
Allow,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum UserPermissionDecision {
AllowOnce,
AlwaysAllow,
Deny,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@ -24,9 +28,9 @@ pub struct Permission {
pub tool_name: String, // "shell", "file_edit", etc.
pub command_pattern: Option<String>, // "rm -rf", "git push", etc.
pub resource_pattern: Option<String>, // "api.search.brave.com", "/src/**", etc.
pub decision: PermissionDecision, // Allow, Deny, Ask
pub scope: PermissionScope, // Session, Project, Global
pub project_id: Option<i32>, // If project-scoped
pub user_decision: UserPermissionDecision,
pub scope: PermissionScope, // Session, Project, Global
pub project_id: Option<i32>, // If project-scoped
pub created_at: DateTime<Utc>,
}
@ -36,7 +40,7 @@ impl Permission {
tool_name: String,
command_pattern: Option<String>,
resource_pattern: Option<String>,
decision: PermissionDecision,
decision: UserPermissionDecision,
scope: PermissionScope,
project_id: Option<i32>,
) -> Self {
@ -45,88 +49,18 @@ impl Permission {
tool_name,
command_pattern,
resource_pattern,
decision,
user_decision: decision,
scope,
project_id,
created_at: Utc::now(),
}
}
/// Check if this permission matches the given tool, command, and path
pub fn matches(&self, tool: &str, command: Option<&str>, path: Option<&PathBuf>) -> bool {
if self.tool_name != tool {
return false;
pub fn system_decision(self) -> SystemPermissionDecision {
if self.user_decision == UserPermissionDecision::AlwaysAllow {
return SystemPermissionDecision::Allow;
}
// Check command pattern if specified
if let (Some(pattern), Some(cmd)) = (&self.command_pattern, command) {
if !self.matches_command(pattern, cmd) {
return false;
}
}
// Check resource pattern if specified
if let (Some(pattern), Some(p)) = (&self.resource_pattern, path) {
if !self.matches_path(pattern, p) {
return false;
}
}
true
}
fn matches_command(&self, pattern: &str, command: &str) -> bool {
// Simple pattern matching - can be enhanced with regex
if pattern.contains('*') {
// Basic wildcard matching
let pattern_parts: Vec<&str> = pattern.split_whitespace().collect();
let command_parts: Vec<&str> = command.split_whitespace().collect();
if pattern_parts.len() != command_parts.len() {
return false;
}
pattern_parts
.iter()
.zip(command_parts.iter())
.all(|(p, c)| *p == "*" || *p == *c)
} else {
pattern == command
}
}
fn matches_path(&self, pattern: &str, path: &PathBuf) -> bool {
let path_str = path.to_string_lossy();
if pattern.contains('*') {
// Basic glob matching
if pattern.ends_with("/**") {
let prefix = pattern.trim_end_matches("/**");
path_str.starts_with(prefix)
} else if pattern.starts_with("**/") {
let suffix = pattern.trim_start_matches("**/");
path_str.ends_with(suffix)
} else {
// Simple wildcard - can be enhanced with proper glob matching
path_str.contains(pattern.trim_matches('*'))
}
} else {
path_str == pattern
}
}
/// Check if this is a session-wide "allow all reads" permission
pub fn is_session_wide_all_reads(&self) -> bool {
matches!(self.decision, PermissionDecision::AllowAllReadsInSession)
&& self.scope == PermissionScope::Session
&& self.resource_pattern.as_ref().map(|p| p.ends_with("/**")).unwrap_or(false)
}
/// Check if this is a session-wide "allow all writes" permission
pub fn is_session_wide_all_writes(&self) -> bool {
matches!(self.decision, PermissionDecision::AllowAllWritesInSession)
&& self.scope == PermissionScope::Session
&& self.resource_pattern.as_ref().map(|p| p.ends_with("/**")).unwrap_or(false)
SystemPermissionDecision::Ask
}
}
@ -165,7 +99,7 @@ impl PermissionRequest {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PermissionConfig {
pub default_decision: PermissionDecision,
pub default_decision: SystemPermissionDecision,
pub dangerous_commands: Vec<String>,
pub restricted_paths: Vec<String>,
pub require_confirmation: bool,
@ -174,7 +108,7 @@ pub struct PermissionConfig {
impl Default for PermissionConfig {
fn default() -> Self {
Self {
default_decision: PermissionDecision::Ask,
default_decision: SystemPermissionDecision::Ask,
dangerous_commands: vec![
"rm -rf".to_string(),
"sudo".to_string(),
@ -188,6 +122,7 @@ impl Default for PermissionConfig {
"/usr/bin".to_string(),
"/bin".to_string(),
"/sbin".to_string(),
".env".to_string(),
"~/.ssh".to_string(),
"~/.aws".to_string(),
"~/.gnupg".to_string(),
@ -195,4 +130,4 @@ impl Default for PermissionConfig {
require_confirmation: true,
}
}
}
}

33
src/domain/prompting/general.rs
View file

@ -1,8 +1,8 @@
use crate::domain::AgentModeType;
use crate::domain::bt::BTStepNodeInterface;
use crate::domain::session::Request;
use crate::domain::ModelType;
use crate::domain::prompting::user::get_user_prompt;
use crate::domain::session::Request;
use crate::domain::AgentModeType;
use crate::domain::ModelType;
/// LLM prompt templates for the coding assistant
///
@ -29,26 +29,36 @@ pub fn get_bt_tree_step_prompt(
}
}
pub fn get_system_prompt(model_type: ModelType, agent_mode: AgentModeType, remaining_calls: usize) -> String {
pub fn get_system_prompt(
model_type: ModelType,
agent_mode: AgentModeType,
remaining_calls: usize,
) -> String {
let (os_name, shell_name) = get_runtime_environment();
if agent_mode == AgentModeType::Plan {
let mut system_prompt = _system_prompt_for_plan(model_type);
if remaining_calls < 3 {
system_prompt.push_str(&format!("\n\nYou have {} tool calls left to process this request.", remaining_calls));
system_prompt.push_str(&format!(
"\n\nYou have {} tool calls left to process this request.",
remaining_calls
));
}
return system_prompt;
}
if agent_mode == AgentModeType::BuildFromPlan {
let mut system_prompt = _system_prompt_for_build_from_plan(model_type);
if remaining_calls < 3 {
system_prompt.push_str(&format!("\n\nYou have {} tool calls left to process this request.", remaining_calls));
system_prompt.push_str(&format!(
"\n\nYou have {} tool calls left to process this request.",
remaining_calls
));
}
return system_prompt;
}
let mut system_prompt = if model_type == ModelType::OpenAI {
format!(
"You are Drastis, a coding agent. \n\
"You're Drastis, you're running as a coding agent in the CLI on a user's computer. \n\
Use the available tools to gather context and make changes. \
When using tools, pass JSON arguments that match their parameters. \n\
Runtime: os={}, shell={}.",
@ -56,12 +66,15 @@ pub fn get_system_prompt(model_type: ModelType, agent_mode: AgentModeType, remai
)
} else {
format!(
"You are Drastis, a coding agent. Use available tools to gather context and make changes. Be concise and accurate. Runtime: os={}, shell={}.",
"You're Drastis, you're running as a coding agent in the CLI on a user's computer. Use available tools to gather context and make changes. Be concise and accurate. Runtime: os={}, shell={}.",
os_name, shell_name
)
};
if remaining_calls < 3 {
system_prompt.push_str(&format!("\n\nYou have {} tool calls left to process this request.", remaining_calls));
system_prompt.push_str(&format!(
"\n\nYou have {} tool calls left to process this request.",
remaining_calls
));
}
system_prompt
}
@ -70,7 +83,7 @@ fn _system_prompt_for_plan(model_type: ModelType) -> String {
let (os_name, shell_name) = get_runtime_environment();
if model_type == ModelType::OpenAI {
format!(
"You are Drastis, a coding agent. \n\
"You're Drastis, you're running as a coding agent in the CLI on a user's computer. \n\
You're in the Plan Mode! Use tools to gather all required information to make a detailed plan of what user wants to achieve. \n\
Ask the user, if there're any ambiguities, and clarify what needs to be done if the instructions are not clear. \n\
When you gathered the information, use the `update_todo_list` tool to create the TODO list and ask the user if he would like to make any changes to the TODO list. \n\

2
src/domain/prompting/mod.rs
View file

@ -4,6 +4,6 @@ mod tools;
mod user;
pub use general::*;
pub use user::*;
pub use session::*;
pub use tools::*;
pub use user::*;

6
src/domain/prompting/user.rs
View file

@ -1,7 +1,7 @@
use crate::domain::AgentModeType;
use crate::domain::ModelType;
use crate::domain::session::Request;
use crate::domain::todo::TodoListStatus;
use crate::domain::AgentModeType;
use crate::domain::ModelType;
/// Format request history and current request as a complete prompt
pub fn get_user_prompt(_model_type: ModelType, request: &dyn Request) -> String {
@ -23,4 +23,4 @@ fn _get_build_from_plan_prompt(request: &dyn Request) -> String {
}
}
"Report very briefly what was done without using any tools, since all the TODO items are completed.".to_string()
}
}

2
src/domain/session/request.rs
View file

@ -41,4 +41,4 @@ pub trait Request {
/// Get the session's TODO list (plan)
fn get_session_plan(&self) -> Option<crate::domain::todo::TodoList>;
}
}

11
src/domain/session/service.rs
View file

@ -7,8 +7,8 @@ use crate::infrastructure::db::DbPool;
use crate::infrastructure::AgentToUiEvent;
use crate::infrastructure::InferenceEngine;
use crate::repository::{
ModelsRepository, SessionRequestStepsRepository, SessionRequestsRepository,
TodoListRepository, UserSettingsRepository,
ModelsRepository, SessionRequestStepsRepository, SessionRequestsRepository, TodoListRepository,
UserSettingsRepository,
};
use crossbeam_channel::Sender;
use std::sync::Arc;
@ -97,7 +97,9 @@ impl SessionService {
match repo.get_by_session(session.id()) {
Ok(Some(row)) => {
match serde_json::from_str::<crate::domain::todo::TodoList>(&row.content) {
Ok(todo_list) => !todo_list.is_completed() && !todo_list.items.is_empty(),
Ok(todo_list) => {
!todo_list.is_completed() && !todo_list.items.is_empty()
}
Err(_) => false,
}
}
@ -133,8 +135,7 @@ impl SessionService {
// Run the workflow
let result: Result<Chain, WorkflowError> = if self.use_behavior_trees {
self.workflow
.run_using_bt(&mut request, cancel)
self.workflow.run_using_bt(&mut request, cancel)
} else {
self.workflow
.run(&mut request, cancel, effective_mode)

36
src/domain/session/session.rs
View file

@ -128,6 +128,7 @@ impl Request for Session {
}
fn get_history_steps(&self) -> Vec<crate::domain::workflow::step::ChainStep> {
use crate::domain::workflow::step::ChainStep;
use crate::repository::SessionRequestStepsRepository;
// BuildFromPlan keeps context minimal — TODO list (with statuses) is
@ -145,16 +146,43 @@ impl Request for Session {
}
};
let steps_repo = SessionRequestStepsRepository::new(conn);
let steps_repo = SessionRequestStepsRepository::new(conn.clone());
// Load all steps from all requests in this session
match steps_repo.load_steps_for_session(self.id) {
let steps = match steps_repo.load_steps_for_session(self.id) {
Ok(steps) => steps,
Err(e) => {
log::warn!("Failed to load history steps for session {}: {}", self.id, e);
log::warn!(
"Failed to load history steps for session {}: {}",
self.id,
e
);
Vec::new()
}
};
if !steps.is_empty() {
return steps;
}
// Fallback: rebuild assistant responses from stored steps_log when steps table is empty
let mut fallback_steps = Vec::new();
for request in &self.requests {
if let Some(log) = request.steps_log() {
for line in log.lines() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
fallback_steps.push(ChainStep::assistant_response(
trimmed.to_string(),
trimmed.to_string(),
));
}
}
}
fallback_steps
}
fn get_session_plan(&self) -> Option<crate::domain::todo::TodoList> {
@ -201,4 +229,4 @@ impl Request for Session {
}
// Note: From<SessionRow> is no longer implemented because Session requires a Project entity.
// Use Session::from_row_with_project() instead.
// Use Session::from_row_with_project() instead.

6
src/domain/session/session_request.rs
View file

@ -8,6 +8,7 @@ use crate::repository::SessionRequestRow;
pub struct SessionRequest {
prompt: String,
result_summary: Option<String>,
steps_log: Option<String>,
#[allow(dead_code)]
mode: AgentModeType,
}
@ -23,6 +24,10 @@ impl SessionRequest {
self.result_summary.as_deref()
}
pub fn steps_log(&self) -> Option<&str> {
self.steps_log.as_deref()
}
#[allow(dead_code)]
pub fn mode(&self) -> AgentModeType {
self.mode
@ -32,6 +37,7 @@ impl SessionRequest {
Self {
prompt: row.prompt,
result_summary: row.result_summary,
steps_log: row._steps_log,
mode: row.mode,
}
}

10
src/domain/startup/service.rs
View file

@ -85,9 +85,10 @@ impl StartupService {
let naming_prompt = session_naming_prompt(self.engine.get_type(), &prompt_preview);
let mut chain = Chain::new();
chain
.steps
.push(ChainStep::user_message(naming_prompt.parse().unwrap(), Vec::new()));
chain.steps.push(ChainStep::user_message(
naming_prompt.parse().unwrap(),
Vec::new(),
));
let session_name = match self.engine.generate(&[], &chain, &[], None) {
Ok(raw) => {
log::debug!("Raw session name response: {:?}", raw.summary);
@ -196,7 +197,8 @@ impl StartupService {
.map(SessionRequest::from_row)
.collect();
let session = Session::load_with_requests(session_row, project, requests);
let mut session = Session::load_with_requests(session_row, project, requests);
session.set_conn(self.conn.clone());
Ok(session)
}
}

18
src/domain/todo.rs
View file

@ -20,9 +20,15 @@ impl TodoListStatus {
#[allow(dead_code)]
pub fn from_str(value: &str) -> Self {
match value {
"pending" | "PENDING" | "planned" | "PLANNED" | "todo" | "TODO" => TodoListStatus::Pending,
"in_progress" | "IN_PROGRESS" | "doing" | "DOING" | "started" | "STARTED" => TodoListStatus::InProgress,
"completed" | "COMPLETED" | "done" | "DONE" | "finished" | "FINISHED" => TodoListStatus::Completed,
"pending" | "PENDING" | "planned" | "PLANNED" | "todo" | "TODO" => {
TodoListStatus::Pending
}
"in_progress" | "IN_PROGRESS" | "doing" | "DOING" | "started" | "STARTED" => {
TodoListStatus::InProgress
}
"completed" | "COMPLETED" | "done" | "DONE" | "finished" | "FINISHED" => {
TodoListStatus::Completed
}
&_ => TodoListStatus::Pending,
}
}
@ -42,6 +48,8 @@ pub struct TodoList {
impl TodoList {
pub fn is_completed(&self) -> bool {
self.items.iter().all(|item| item.status == TodoListStatus::Completed)
self.items
.iter()
.all(|item| item.status == TodoListStatus::Completed)
}
}
}

14
src/domain/tools/discover_objects.rs
View file

@ -106,7 +106,12 @@ impl Tool for DiscoverObjects {
let input = match self.load_input() {
Ok(input) => input,
Err(e) => {
return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new())
return ToolResult::error(
self.name().to_string(),
String::new(),
e.to_string(),
String::new(),
)
}
};
@ -117,7 +122,12 @@ impl Tool for DiscoverObjects {
Self::format_output(lang, &objects),
input.call_id,
),
Err(e) => ToolResult::error(self.name().to_string(), input.raw.clone(), e.to_string(), input.call_id),
Err(e) => ToolResult::error(
self.name().to_string(),
input.raw.clone(),
e.to_string(),
input.call_id,
),
}
}

16
src/domain/tools/find_files.rs
View file

@ -169,7 +169,12 @@ impl Tool for FindFiles {
let input = match self.load_input() {
Ok(input) => input,
Err(e) => {
return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new())
return ToolResult::error(
self.name().to_string(),
String::new(),
e.to_string(),
String::new(),
)
}
};
@ -177,7 +182,14 @@ impl Tool for FindFiles {
let search_root =
match Self::resolve_search_root(input.root.as_deref(), request.project_root()) {
Ok(root) => root,
Err(e) => return ToolResult::error(self.name().to_string(), input.raw.clone(), e, input.call_id.clone()),
Err(e) => {
return ToolResult::error(
self.name().to_string(),
input.raw.clone(),
e,
input.call_id.clone(),
)
}
};
// Check if search root exists

7
src/domain/tools/mod.rs
View file

@ -47,7 +47,7 @@ pub struct ToolResult {
tool_name: String,
call_id: String,
pub(crate) input: String,
is_successful: bool,
is_successful: bool,
output: String,
error_message: String,
file_changes: Option<Vec<FileChange>>,
@ -62,7 +62,7 @@ impl ToolResult {
output,
error_message: String::new(),
file_changes: None,
call_id
call_id,
}
}
@ -74,7 +74,7 @@ impl ToolResult {
output: String::new(),
error_message: message.into(),
file_changes: None,
call_id
call_id,
}
}
@ -161,6 +161,7 @@ pub trait Tool: Send + Sync {
fn is_read_only(&self) -> bool {
false
}
#[allow(dead_code)]
fn skip_permission_check(&self) -> bool {
false
}

11
src/domain/tools/patch_files.rs
View file

@ -78,7 +78,12 @@ impl Tool for PatchFiles {
let input = match self.load_input() {
Ok(input) => input,
Err(e) => {
return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new())
return ToolResult::error(
self.name().to_string(),
String::new(),
e.to_string(),
String::new(),
)
}
};
@ -149,7 +154,7 @@ impl Tool for PatchFiles {
return ToolResult::error(
self.name().to_string(),
input.raw.clone(),
format!("Failed to apply patch: {}", e),
format!("Patch Failed: {}. Next step: read the file region you’re editing (function/block + ~10 lines around it) and generate a new patch using those exact lines as context. Then retry patch_files.", e),
input.call_id.clone(),
)
}
@ -612,4 +617,4 @@ mod tests {
result.output_string()
);
}
}
}

16
src/domain/tools/read_objects.rs
View file

@ -202,7 +202,12 @@ impl Tool for ReadObjects {
let input = match self.load_input() {
Ok(input) => input,
Err(e) => {
return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new())
return ToolResult::error(
self.name().to_string(),
String::new(),
e.to_string(),
String::new(),
)
}
};
@ -213,7 +218,12 @@ impl Tool for ReadObjects {
Self::format_output(lang, results),
input.call_id,
),
Err(e) => ToolResult::error(self.name().to_string(), input.raw.clone(), e.to_string(), input.call_id),
Err(e) => ToolResult::error(
self.name().to_string(),
input.raw.clone(),
e.to_string(),
input.call_id,
),
}
}
@ -485,4 +495,4 @@ mod tests {
};
assert!(query.matches(&obj_partial), "Should match partial name");
}
}
}

147
src/domain/tools/shell_exec.rs
View file

@ -28,6 +28,35 @@ struct ShellExecInputParsed {
call_id: String,
}
fn is_allowed_read_only_command(command: &str) -> bool {
let trimmed = command.trim();
if trimmed.is_empty() {
return false;
}
if trimmed.contains(';')
|| trimmed.contains("&&")
|| trimmed.contains("||")
|| trimmed.contains('|')
|| trimmed.contains('>')
|| trimmed.contains('<')
{
return false;
}
let mut parts = trimmed.split_whitespace();
let first = parts.next().unwrap_or_default();
let args: Vec<&str> = parts.collect();
match first {
"rg" | "grep" | "glob" | "cat" | "head" | "tail" | "less" | "more" | "wc" | "cut"
| "sort" | "uniq" | "find" | "ls" | "tree" | "stat" | "file" | "awk" | "pwd"
| "which" | "type" => true,
"sed" => args.iter().any(|arg| *arg == "-n" || *arg == "--quiet" || *arg == "--silent"),
_ => false,
}
}
impl Tool for ShellExec {
fn name(&self) -> &'static str {
"shell_exec"
@ -135,7 +164,12 @@ impl Tool for ShellExec {
if !stderr.is_empty() {
result.push_str(&format!("[stderr]: {}", stderr));
}
return ToolResult::error(self.name().to_string(), input.raw.clone(), result, input.call_id.clone());
return ToolResult::error(
self.name().to_string(),
input.raw.clone(),
result,
input.call_id.clone(),
);
};
ToolResult::ok(self.name().to_string(), input.raw, result, input.call_id)
@ -251,6 +285,15 @@ DO NOT use it to change files, use `patch_files` tool for this.",
paths
}
fn is_read_only(&self) -> bool {
self.input
.lock()
.unwrap()
.as_ref()
.map(|input| is_allowed_read_only_command(&input.command))
.unwrap_or(false)
}
}
impl ShellExec {
@ -346,7 +389,10 @@ mod tests {
let tool = ShellExec::new();
assert!(tool
.parse_input(r#"{"command":"echo hello"}"#.to_string(), "call-id".to_string())
.parse_input(
r#"{"command":"echo hello"}"#.to_string(),
"call-id".to_string()
)
.is_none());
let result = tool.work(&request);
@ -388,10 +434,13 @@ mod tests {
let tool = ShellExec::new();
assert!(tool
.parse_input(format!(
r#"{{"command":"pwd","working_dir":"{}"}}"#,
subdir.display()
), "call-id".to_string())
.parse_input(
format!(
r#"{{"command":"pwd","working_dir":"{}"}}"#,
subdir.display()
),
"call-id".to_string()
)
.is_none());
let result = tool.work(&request);
@ -412,7 +461,10 @@ mod tests {
let tool = ShellExec::new();
assert!(tool
.parse_input(r#"{"command":"pwd","working_dir":"/tmp"}"#.to_string(), "call-id".to_string())
.parse_input(
r#"{"command":"pwd","working_dir":"/tmp"}"#.to_string(),
"call-id".to_string()
)
.is_none());
let result = tool.work(&request);
@ -450,7 +502,10 @@ mod tests {
let tool = ShellExec::new();
assert!(tool
.parse_input(r#"{"command":"echo error >&2 && exit 1"}"#.to_string(), "call-id".to_string())
.parse_input(
r#"{"command":"echo error >&2 && exit 1"}"#.to_string(),
"call-id".to_string()
)
.is_none());
let result = tool.work(&request);
@ -478,10 +533,13 @@ mod tests {
let tool = ShellExec::new();
assert!(tool
.parse_input(format!(
r#"{{"command":"echo 'test content' > {}"}}"#,
file_path.display()
), "call-id".to_string())
.parse_input(
format!(
r#"{{"command":"echo 'test content' > {}"}}"#,
file_path.display()
),
"call-id".to_string()
)
.is_none());
let result = tool.work(&request);
@ -503,7 +561,10 @@ mod tests {
let tool = ShellExec::new();
assert!(tool
.parse_input(r#"{"command":"echo 'hello world' | tr 'a-z' 'A-Z'"}"#.to_string(), "call-id".to_string())
.parse_input(
r#"{"command":"echo 'hello world' | tr 'a-z' 'A-Z'"}"#.to_string(),
"call-id".to_string()
)
.is_none());
let result = tool.work(&request);
@ -513,4 +574,64 @@ mod tests {
result.output_string()
);
}
#[test]
fn test_shell_exec_is_read_only_allows_read_commands() {
let tool = ShellExec::new();
assert!(tool
.parse_input(
r#"{"command":"rg -n foo src"}"#.to_string(),
"call-id".to_string()
)
.is_none());
assert!(tool.is_read_only());
let tool = ShellExec::new();
assert!(tool
.parse_input(
r#"{"command":"grep -n foo file.txt"}"#.to_string(),
"call-id".to_string()
)
.is_none());
assert!(tool.is_read_only());
let tool = ShellExec::new();
assert!(tool
.parse_input(
r#"{"command":"sed -n '1,10p' file.txt"}"#.to_string(),
"call-id".to_string()
)
.is_none());
assert!(tool.is_read_only());
}
#[test]
fn test_shell_exec_is_read_only_rejects_write_or_compound() {
let tool = ShellExec::new();
assert!(tool
.parse_input(
r#"{"command":"sed '1,10p' file.txt"}"#.to_string(),
"call-id".to_string()
)
.is_none());
assert!(!tool.is_read_only());
let tool = ShellExec::new();
assert!(tool
.parse_input(
r#"{"command":"echo hi > out.txt"}"#.to_string(),
"call-id".to_string()
)
.is_none());
assert!(!tool.is_read_only());
let tool = ShellExec::new();
assert!(tool
.parse_input(
r#"{"command":"echo hi && rg -n foo src"}"#.to_string(),
"call-id".to_string()
)
.is_none());
assert!(!tool.is_read_only());
}
}

16
src/domain/tools/structure.rs
View file

@ -164,13 +164,25 @@ impl Tool for Structure {
let input = match self.load_input() {
Ok(input) => input,
Err(e) => {
return ToolResult::error(self.name().to_string(), String::new(), e.to_string(), String::new())
return ToolResult::error(
self.name().to_string(),
String::new(),
e.to_string(),
String::new(),
)
}
};
let target_path = match Self::resolve_path(&input.path, request.project_root()) {
Ok(p) => p,
Err(e) => return ToolResult::error(self.name().to_string(), input.raw.clone(), e, input.call_id.clone()),
Err(e) => {
return ToolResult::error(
self.name().to_string(),
input.raw.clone(),
e,
input.call_id.clone(),
)
}
};
if !target_path.exists() {

9
src/domain/tools/update_todo_list.rs
View file

@ -1,4 +1,5 @@
use crate::domain::session::Request;
use crate::domain::todo::TodoList;
use crate::domain::tools::{Error, Tool, ToolResult};
use crate::infrastructure::db::DbPool;
use crate::infrastructure::event_bus::AgentToUiEvent;
@ -6,7 +7,6 @@ use crate::repository::TodoListRepository;
use crossbeam_channel::Sender;
use serde_json::{json, Value};
use std::sync::Mutex;
use crate::domain::todo::TodoList;
pub struct UpdateTodoList {
input: Mutex<Option<UpdateTodoListInput>>,
@ -42,7 +42,10 @@ impl Tool for UpdateTodoList {
match parsed {
Ok(_) => {
let mut lock = self.input.lock().unwrap();
*lock = Some(UpdateTodoListInput { raw: input, call_id });
*lock = Some(UpdateTodoListInput {
raw: input,
call_id,
});
None
}
Err(e) => Some(Error::Parse(format!("Failed to parse input: {}", e))),
@ -148,7 +151,7 @@ impl Tool for UpdateTodoList {
}
fn desc(&self) -> String {
"Use this in Plan mode to create or update TODO items after exploring the codebase. For every use, provide full TODO list, as it will replace the existing list.\n\
"Use this in Plan mode to create or update TODO items after exploring the codebase. For every use, provide full TODO list, as it will replace the existing list.\n\
If you're not in Plan mode, only update statuses of TODO items as you finish them.
"
.to_string()

43
src/domain/tools/web_search.rs
View file

@ -163,9 +163,12 @@ impl Tool for WebSearch {
Ok(results) => {
let output = WebSearchOutput { results };
match serde_json::to_string(&output) {
Ok(json_output) => {
ToolResult::ok(self.name().to_string(), input.raw.clone(), json_output, input.call_id)
}
Ok(json_output) => ToolResult::ok(
self.name().to_string(),
input.raw.clone(),
json_output,
input.call_id,
),
Err(e) => ToolResult::error(
self.name().to_string(),
input.raw.clone(),
@ -320,7 +323,10 @@ mod tests {
use crate::domain::permissions::{
checker::{PermissionChecker, PermissionPrompter},
store::{PermissionStore, StoreError},
types::{PermissionConfig, PermissionDecision, PermissionRequest, PermissionScope},
types::{
PermissionConfig, PermissionRequest, PermissionScope, SystemPermissionDecision,
UserPermissionDecision,
},
};
use crate::domain::session::{Request, SessionRequest};
use crate::domain::UserSettings;
@ -359,11 +365,6 @@ mod tests {
Ok(permission)
}
fn find_session_permissions(&self, _project_id: i32) -> Result<Vec<crate::domain::permissions::types::Permission>, StoreError> {
// Return empty for tests - this test doesn't use session permissions
Ok(Vec::new())
}
fn find_permission(
&self,
tool: &str,
@ -407,14 +408,14 @@ mod tests {
struct TestPrompter {
calls: Arc<AtomicUsize>,
decision: PermissionDecision,
decision: UserPermissionDecision,
}
impl PermissionPrompter for TestPrompter {
fn ask_permission(
&self,
_request: &PermissionRequest,
) -> Result<PermissionDecision, crate::utils::AskError> {
) -> Result<UserPermissionDecision, crate::utils::AskError> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(self.decision.clone())
}
@ -482,6 +483,11 @@ mod tests {
web_search_enabled,
brave_api_key: brave_api_key.map(String::from),
max_tool_calls_per_request: 10,
auth_method: "api_key".to_string(),
oauth_access_token: None,
oauth_refresh_token: None,
oauth_token_expiry: None,
oauth_account_id: None,
};
UserSettings::from(row)
}
@ -495,12 +501,12 @@ mod tests {
let calls = Arc::new(AtomicUsize::new(0));
let prompter = Arc::new(TestPrompter {
calls: Arc::clone(&calls),
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AllowOnce,
});
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::Ask,
default_decision: SystemPermissionDecision::Ask,
..PermissionConfig::default()
},
prompter,
@ -525,12 +531,12 @@ mod tests {
let store = Arc::new(TestStore::new());
let calls = Arc::new(AtomicUsize::new(0));
let decision_to_return = Arc::new(Mutex::new(PermissionDecision::AllowOnce));
let decision_to_return = Arc::new(Mutex::new(UserPermissionDecision::AllowOnce));
let captured_request = Arc::new(Mutex::new(None::<PermissionRequest>));
struct CapturingPrompter {
calls: Arc<AtomicUsize>,
decision: Arc<Mutex<PermissionDecision>>,
decision: Arc<Mutex<UserPermissionDecision>>,
captured: Arc<Mutex<Option<PermissionRequest>>>,
}
@ -538,7 +544,7 @@ mod tests {
fn ask_permission(
&self,
request: &PermissionRequest,
) -> Result<PermissionDecision, crate::utils::AskError> {
) -> Result<UserPermissionDecision, crate::utils::AskError> {
self.calls.fetch_add(1, Ordering::SeqCst);
*self.captured.lock().unwrap() = Some(request.clone());
Ok(self.decision.lock().unwrap().clone())
@ -554,7 +560,7 @@ mod tests {
let checker = PermissionChecker::new_with_prompter(
store,
PermissionConfig {
default_decision: PermissionDecision::Ask,
default_decision: SystemPermissionDecision::Ask,
..PermissionConfig::default()
},
prompter,
@ -578,5 +584,4 @@ mod tests {
assert_eq!(captured_req.tool_name, "web_search");
assert_eq!(captured_req.paths, vec![PathBuf::from("docs.rs")]);
}
}
}

59
src/domain/user_settings.rs
View file

@ -1,6 +1,30 @@
use crate::repository::UserSettingsRow;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthMethod {
ApiKey,
OAuth,
}
#[allow(dead_code)]
impl AuthMethod {
pub fn as_str(&self) -> &str {
match self {
AuthMethod::ApiKey => "api_key",
AuthMethod::OAuth => "oauth",
}
}
pub fn from_str(s: &str) -> Self {
match s {
"oauth" => AuthMethod::OAuth,
_ => AuthMethod::ApiKey,
}
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct UserSettings {
_id: i64,
openai_api_key: Option<String>,
@ -11,8 +35,14 @@ pub struct UserSettings {
web_search_enabled: bool,
brave_api_key: Option<String>,
max_tool_calls_per_request: i32,
auth_method: AuthMethod,
oauth_access_token: Option<String>,
oauth_refresh_token: Option<String>,
oauth_token_expiry: Option<i64>,
oauth_account_id: Option<String>,
}
#[allow(dead_code)]
impl UserSettings {
pub fn openai_api_key(&self) -> Option<&str> {
self.openai_api_key.as_deref()
@ -22,7 +52,6 @@ impl UserSettings {
self.openai_tracing_enabled
}
#[allow(dead_code)]
pub fn current_model_name(&self) -> Option<&str> {
self.current_model_name.as_deref()
}
@ -43,6 +72,29 @@ impl UserSettings {
pub fn max_tool_calls_per_request(&self) -> i32 {
self.max_tool_calls_per_request
}
pub fn auth_method(&self) -> &AuthMethod {
&self.auth_method
}
pub fn oauth_access_token(&self) -> Option<&str> {
self.oauth_access_token.as_deref()
}
pub fn oauth_refresh_token(&self) -> Option<&str> {
self.oauth_refresh_token.as_deref()
}
pub fn oauth_account_id(&self) -> Option<&str> {
self.oauth_account_id.as_deref()
}
pub fn is_oauth_token_expired(&self) -> bool {
if let Some(expiry) = self.oauth_token_expiry {
return chrono::Utc::now().timestamp() >= expiry;
}
false
}
}
impl From<UserSettingsRow> for UserSettings {
@ -57,6 +109,11 @@ impl From<UserSettingsRow> for UserSettings {
web_search_enabled: row.web_search_enabled,
brave_api_key: row.brave_api_key,
max_tool_calls_per_request: row.max_tool_calls_per_request,
auth_method: AuthMethod::from_str(&row.auth_method),
oauth_access_token: row.oauth_access_token,
oauth_refresh_token: row.oauth_refresh_token,
oauth_token_expiry: row.oauth_token_expiry,
oauth_account_id: row.oauth_account_id,
}
}
}

9
src/domain/workflow/chain.rs
View file

@ -1,8 +1,8 @@
use super::step::ChainStep;
use super::step::StepType;
use crate::domain::todo::TodoList;
use crate::domain::tools::ToolResult;
use serde::{Deserialize, Serialize};
use crate::domain::todo::TodoList;
use super::step::StepType;
/// Represents an execution chain containing multiple steps
/// The chain is built incrementally as tools are executed
@ -36,7 +36,8 @@ impl Chain {
/// Add a step to the chain after executing a tool
pub fn add_step(&mut self, result: ToolResult) {
if result.tool_name() == "update_todo_list" && result.is_successful() {
if let Ok(updated_todo_list) = serde_json::from_str::<TodoList>(&result.input_string()) {
if let Ok(updated_todo_list) = serde_json::from_str::<TodoList>(&result.input_string())
{
self.set_todo_list(Some(updated_todo_list));
return;
}
@ -122,7 +123,7 @@ impl Chain {
.collect::<Vec<_>>()
.join("\n")
}
pub fn set_system_prompt(&mut self, system_prompt: String) {
self.system_prompt = system_prompt;
}

2
src/domain/workflow/toolset/all.rs
View file

@ -16,7 +16,7 @@ impl AllToolset {
session_id: i64,
settings: &UserSettings,
conn: DbPool,
event_sender: Sender<AgentToUiEvent>
event_sender: Sender<AgentToUiEvent>,
) -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = HashMap::new();

4
src/domain/workflow/toolset/all_no_todo.rs
View file

@ -15,7 +15,7 @@ impl AllNoTodoToolset {
_session_id: i64,
settings: &UserSettings,
_conn: DbPool,
_event_sender: Sender<AgentToUiEvent>
_event_sender: Sender<AgentToUiEvent>,
) -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = HashMap::new();
@ -50,4 +50,4 @@ impl Toolset for AllNoTodoToolset {
fn tools(&self) -> &HashMap<String, Box<dyn Tool>> {
&self.tools
}
}
}

4
src/domain/workflow/toolset/discover.rs
View file

@ -1,9 +1,9 @@
use crate::domain::tools::*;
use crate::domain::workflow::toolset::Toolset;
use crate::infrastructure::db::DbPool;
use std::collections::HashMap;
use crossbeam_channel::Sender;
use crate::infrastructure::AgentToUiEvent;
use crossbeam_channel::Sender;
use std::collections::HashMap;
/// General toolset containing read-only and utility tools
pub struct DiscoverToolset {

6
src/domain/workflow/toolset/finishing_all.rs
View file

@ -11,11 +11,7 @@ pub struct FinishingAllToolset {
}
impl FinishingAllToolset {
pub fn new(
session_id: i64,
conn: DbPool,
event_sender: Sender<AgentToUiEvent>
) -> Self {
pub fn new(session_id: i64, conn: DbPool, event_sender: Sender<AgentToUiEvent>) -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = HashMap::new();
let patch_files = Box::new(PatchFiles::new());

8
src/domain/workflow/toolset/finishing_all_no_todo.rs
View file

@ -10,11 +10,7 @@ pub struct FinishingAllNoTodoToolset {
}
impl FinishingAllNoTodoToolset {
pub fn new(
_session_id: i64,
_conn: DbPool,
_event_sender: Sender<AgentToUiEvent>
) -> Self {
pub fn new(_session_id: i64, _conn: DbPool, _event_sender: Sender<AgentToUiEvent>) -> Self {
let mut tools: HashMap<String, Box<dyn Tool>> = HashMap::new();
let patch_files = Box::new(PatchFiles::new());
@ -28,4 +24,4 @@ impl Toolset for FinishingAllNoTodoToolset {
fn tools(&self) -> &HashMap<String, Box<dyn Tool>> {
&self.tools
}
}
}

58
src/domain/workflow/toolset/mod.rs
View file

@ -1,29 +1,29 @@
mod all;
mod all_no_todo;
mod edit;
mod none;
mod discover;
mod edit;
mod finishing_all;
mod finishing_all_no_todo;
mod finishing_discover;
mod none;
use std::collections::HashMap;
use std::sync::Arc;
use crossbeam_channel::Sender;
use crate::domain::tools::Error;
use crate::domain::tools::Tool;
use crate::domain::UserSettings;
use crate::infrastructure::db::DbPool;
use crate::infrastructure::event_bus::AgentToUiEvent;
use crate::infrastructure::inference::ToolCall;
pub use all::AllToolset;
pub use all_no_todo::AllNoTodoToolset;
pub use edit::EditToolset;
pub use none::NoneToolset;
use crossbeam_channel::Sender;
pub use discover::DiscoverToolset;
pub use edit::EditToolset;
pub use finishing_all::FinishingAllToolset;
pub use finishing_all_no_todo::FinishingAllNoTodoToolset;
pub use finishing_discover::FinishingDiscoverToolset;
use crate::infrastructure::db::DbPool;
use crate::domain::tools::Tool;
use crate::domain::UserSettings;
use crate::infrastructure::event_bus::AgentToUiEvent;
use crate::infrastructure::inference::ToolCall;
use crate::domain::tools::Error;
pub use none::NoneToolset;
use std::collections::HashMap;
use std::sync::Arc;
#[derive(Clone, Copy)]
pub enum ToolsetType {
@ -51,12 +51,29 @@ impl ToolsetType {
ToolsetType::Discover => Arc::new(DiscoverToolset::new(session_id, conn, event_sender)),
ToolsetType::Edit => Arc::new(EditToolset::new(session_id, conn, event_sender)),
ToolsetType::All => Arc::new(AllToolset::new(session_id, settings, conn, event_sender)),
ToolsetType::AllNoTodo => Arc::new(AllNoTodoToolset::new(session_id, settings, conn, event_sender)),
ToolsetType::FinishingAllNoTodo => Arc::new(FinishingAllNoTodoToolset::new(session_id, conn, event_sender)),
ToolsetType::AllNoTodo => Arc::new(AllNoTodoToolset::new(
session_id,
settings,
conn,
event_sender,
)),
ToolsetType::FinishingAllNoTodo => Arc::new(FinishingAllNoTodoToolset::new(
session_id,
conn,
event_sender,
)),
ToolsetType::None => Arc::new(NoneToolset::new()),
ToolsetType::FinishingAll => Arc::new(FinishingAllToolset::new(session_id, conn, event_sender)),
ToolsetType::FinishingDiscover => Arc::new(FinishingDiscoverToolset::new(session_id, conn, event_sender)),
ToolsetType::FinishingEdit => Arc::new(EditToolset::new(session_id, conn, event_sender)),
ToolsetType::FinishingAll => {
Arc::new(FinishingAllToolset::new(session_id, conn, event_sender))
}
ToolsetType::FinishingDiscover => Arc::new(FinishingDiscoverToolset::new(
session_id,
conn,
event_sender,
)),
ToolsetType::FinishingEdit => {
Arc::new(EditToolset::new(session_id, conn, event_sender))
}
}
}
@ -102,10 +119,11 @@ pub trait Toolset {
.map(|t| t.as_ref())
.ok_or_else(|| Error::Parse(format!("Tool not found: {}", tool_call.name)))?;
if let Some(err) = tool.parse_input(tool_call.arguments.clone(), tool_call.call_id.clone()) {
if let Some(err) = tool.parse_input(tool_call.arguments.clone(), tool_call.call_id.clone())
{
return Err(err);
}
Ok(tool)
}
}
}

62
src/domain/workflow/workflow.rs
View file

@ -6,17 +6,17 @@ use crate::domain::bt::GeneralTree;
use crate::domain::prompting;
use crate::domain::session::Request;
use crate::domain::todo::TodoListStatus;
use crate::domain::workflow::toolset::NoneToolset;
use crate::domain::workflow::toolset::ToolsetType;
use crate::domain::AgentModeType;
use crate::infrastructure::db::DbPool;
use crate::infrastructure::event_bus::{AgentToUiEvent, StepPhase};
use crate::infrastructure::inference::InferenceEngine;
use tokio::runtime::{Builder, Handle};
use crate::domain::workflow::toolset::NoneToolset;
use crate::domain::workflow::toolset::ToolsetType;
use crossbeam_channel::Sender;
use openai_agents_tracing::{SpanKind, TracingFacade};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use crossbeam_channel::Sender;
use crate::domain::AgentModeType;
use tokio::runtime::{Builder, Handle};
/// Main workflow orchestrator that runs LLM-driven coding tasks
/// Implements an eternal agent loop that:
@ -48,10 +48,7 @@ impl Workflow {
toolset: Arc::new(NoneToolset::new()),
engine,
chain: Chain::new(),
tool_runner: ToolRunner::new(
permission_checker,
event_sender.clone(),
),
tool_runner: ToolRunner::new(permission_checker, event_sender.clone()),
tracer: None,
event_sender,
conn,
@ -170,17 +167,20 @@ impl Workflow {
user_prompt_override: Option<String>,
mode: AgentModeType,
) -> Result<(), Error> {
let base_user_prompt = prompting::get_user_prompt(self.engine.get_type(), request);
let user_prompt = user_prompt_override
.as_deref()
.unwrap_or(&base_user_prompt)
.to_string();
self.chain.steps.push(ChainStep::user_message(user_prompt.clone(), request.images().to_vec()));
self.chain.steps.push(ChainStep::user_message(
user_prompt.clone(),
request.images().to_vec(),
));
// Get max_tool_calls from override (for BT mode) or user settings
let max_tool_calls = max_tool_calls_override.unwrap_or_else(|| {
request.user_settings()
request
.user_settings()
.map(|s| s.max_tool_calls_per_request() as usize)
.unwrap_or(50)
});
@ -189,7 +189,11 @@ impl Workflow {
let mut tool_call_count = 0;
let mut current_active_todo = self.chain.todo_list.clone();
let mut in_finishing_mode = false;
let finishing_threshold = if max_tool_calls > 5 { max_tool_calls - 5 } else { max_tool_calls };
let finishing_threshold = if max_tool_calls > 5 {
max_tool_calls - 5
} else {
max_tool_calls
};
// Eternal agent loop
for _iteration in 1..=max_tool_calls {
@ -204,7 +208,11 @@ impl Workflow {
let remaining_calls = max_tool_calls.saturating_sub(tool_call_count);
// Get base system prompt and inject remaining count
let system_prompt = prompting::get_system_prompt(self.engine.get_type(), request.mode(), remaining_calls);
let system_prompt = prompting::get_system_prompt(
self.engine.get_type(),
request.mode(),
remaining_calls,
);
self.chain.set_system_prompt(system_prompt.clone());
// Switch to finishing toolset if approaching limit
@ -280,7 +288,10 @@ impl Workflow {
// Check if TODO item changed and reset counter if so (BuildFromPlan mode only)
if mode == AgentModeType::BuildFromPlan && is_update_todo {
if self._did_todo_item_change(&current_active_todo) {
log::info!("Active TODO item changed, resetting counter from {} to 0", tool_call_count);
log::info!(
"Active TODO item changed, resetting counter from {} to 0",
tool_call_count
);
tool_call_count = 0;
current_active_todo = self.chain.todo_list.clone();
@ -364,29 +375,28 @@ impl Workflow {
.map(|duration| duration.as_nanos())
.unwrap_or(0);
let index = (nanos % options.len() as u128) as usize;
let _ = self
.event_sender
.send(AgentToUiEvent::ProgressEvent {
step_name: "inference".to_string(),
phase: StepPhase::Before,
summary: options[index].to_string(),
});
let _ = self.event_sender.send(AgentToUiEvent::ProgressEvent {
step_name: "inference".to_string(),
phase: StepPhase::Before,
summary: options[index].to_string(),
});
}
/// Detects if the active TODO item has changed
/// Returns true if the first non-completed item's title is different
fn _did_todo_item_change(&self, previous_todo: &Option<crate::domain::todo::TodoList>) -> bool {
// Get first non-completed item from previous TODO list
let prev_active = previous_todo.as_ref().and_then(|list| {
list.items.iter()
list.items
.iter()
.find(|item| item.status != TodoListStatus::Completed)
.map(|item| &item.title)
});
// Get first non-completed item from current TODO list
let curr_active = self.chain.todo_list.as_ref().and_then(|list| {
list.items.iter()
list.items
.iter()
.find(|item| item.status != TodoListStatus::Completed)
.map(|item| &item.title)
});
@ -396,4 +406,4 @@ impl Workflow {
// If both exist but titles differ, change occurred
prev_active != curr_active
}
}
}

387
src/infrastructure/api_clients/openai/client.rs
View file

@ -3,6 +3,15 @@ use super::translator::{build_llm_result, build_request_dto};
use crate::infrastructure::inference::LLMInferenceResult;
use std::error::Error;
#[derive(Clone)]
pub enum AuthToken {
ApiKey(String),
OAuth {
token: String,
account_id: Option<String>,
},
}
#[derive(Debug)]
pub enum OpenAIClientError {
Api {
@ -16,6 +25,17 @@ pub enum OpenAIClientError {
NoText {
body: String,
},
BodyRead {
source: reqwest::Error,
status: reqwest::StatusCode,
content_encoding: Option<String>,
content_type: Option<String>,
content_length: Option<u64>,
transfer_encoding: Option<String>,
response_headers: String,
partial_body: Option<String>,
},
AuthExpired,
}
impl std::fmt::Display for OpenAIClientError {
@ -34,6 +54,34 @@ impl std::fmt::Display for OpenAIClientError {
OpenAIClientError::NoText { body } => {
write!(f, "No output_text found in OpenAI response (body={})", body)
}
OpenAIClientError::BodyRead {
source,
status,
content_encoding,
content_type,
content_length,
transfer_encoding,
response_headers,
partial_body,
} => {
write!(
f,
"Failed to read SSE response body (status={}, content-encoding={}, content-type={}, content-length={}, transfer-encoding={}, headers={}, partial-body={}, error={:?})",
status,
content_encoding.as_deref().unwrap_or("<missing>"),
content_type.as_deref().unwrap_or("<missing>"),
content_length
.map(|value| value.to_string())
.unwrap_or_else(|| "<missing>".to_string()),
transfer_encoding.as_deref().unwrap_or("<missing>"),
response_headers,
partial_body.as_deref().unwrap_or("<empty>"),
source
)
}
OpenAIClientError::AuthExpired => {
write!(f, "OAuth token expired - please re-authenticate")
}
}
}
}
@ -42,23 +90,27 @@ impl std::error::Error for OpenAIClientError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
OpenAIClientError::Deserialize { source, .. } => Some(source),
OpenAIClientError::BodyRead { source, .. } => Some(source),
_ => None,
}
}
}
pub struct OpenAIClient {
api_key: String,
auth_token: AuthToken,
model: String,
client: reqwest::blocking::Client,
}
impl OpenAIClient {
pub fn new(api_key: String, model: String) -> Self {
pub fn new(auth_token: AuthToken, model: String) -> Self {
let client = reqwest::blocking::Client::builder()
.build()
.expect("Failed to build OpenAI HTTP client");
Self {
api_key,
auth_token,
model,
client: reqwest::blocking::Client::new(),
client,
}
}
@ -71,12 +123,7 @@ impl OpenAIClient {
) -> Result<LLMInferenceResult, Box<dyn Error + Send + Sync>> {
let max_attempts = 3;
for attempt in 1..=max_attempts {
match self.call_responses_api_inner(
tools,
chain,
images,
tracer.as_deref_mut(),
) {
match self.call_responses_api_inner(tools, chain, images, tracer.as_deref_mut()) {
Ok(result) => return Ok(result),
Err(e) => {
let error_str = e.to_string();
@ -122,15 +169,14 @@ impl OpenAIClient {
images: &[String],
mut tracer: Option<&mut openai_agents_tracing::TracingFacade>,
) -> Result<LLMInferenceResult, Box<dyn Error + Send + Sync>> {
let url = "https://api.openai.com/v1/responses";
// Determine API URL based on auth type
let url = match &self.auth_token {
AuthToken::OAuth { .. } => "https://chatgpt.com/backend-api/codex/responses",
AuthToken::ApiKey(_) => "https://api.openai.com/v1/responses",
};
let request_body = build_request_dto(
&self.model,
images,
tools,
chain,
tracer.as_deref_mut(),
);
let request_body =
build_request_dto(&self.model, images, tools, chain, tracer.as_deref_mut());
// Start span with model name and add request as JSON
if let Some(tracer) = &mut tracer {
@ -142,30 +188,97 @@ impl OpenAIClient {
}
}
let response = self
.client
.post(url)
.header("Authorization", format!("Bearer {}", self.api_key))
.json(&request_body)
// Build request with appropriate headers
let mut request_builder = self.client.post(url).json(&request_body);
// Add authentication header
match &self.auth_token {
AuthToken::ApiKey(key) => {
request_builder =
request_builder.header("Authorization", format!("Bearer {}", key));
}
AuthToken::OAuth { token, account_id } => {
request_builder =
request_builder.header("Authorization", format!("Bearer {}", token));
if let Some(acc_id) = account_id {
request_builder = request_builder.header("ChatGPT-Account-Id", acc_id);
}
}
}
let mut response = request_builder
.header("Accept", "text/event-stream")
.header(reqwest::header::ACCEPT_ENCODING, "identity")
.send()?;
let status = response.status();
let body = response.text()?;
let content_encoding = response
.headers()
.get(reqwest::header::CONTENT_ENCODING)
.and_then(|value| value.to_str().ok())
.map(|value| value.to_string());
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(|value| value.to_string());
let content_length = response.content_length();
let transfer_encoding = response
.headers()
.get(reqwest::header::TRANSFER_ENCODING)
.and_then(|value| value.to_str().ok())
.map(|value| value.to_string());
let response_headers = Self::format_headers_excerpt(response.headers());
if !status.is_success() {
let body = response
.text()
.unwrap_or_else(|e| format!("(failed to read body: {})", e));
if let Some(t) = tracer {
t.end_span(&self.model);
}
// Check for OAuth expiration
if status.as_u16() == 401 {
return Err(Box::new(OpenAIClientError::AuthExpired));
}
return Err(Box::new(OpenAIClientError::Api { status, body }));
}
let dto = match serde_json::from_str::<ResponseDTO>(&body) {
// Read SSE stream as raw bytes to avoid encoding issues with chunked transfer
let mut body_bytes = Vec::new();
if let Err(e) = response.copy_to(&mut body_bytes) {
let partial_body = if body_bytes.is_empty() {
None
} else {
Some(Self::format_body_excerpt(
&String::from_utf8_lossy(&body_bytes),
4000,
))
};
return Err(Box::new(OpenAIClientError::BodyRead {
source: e,
status,
content_encoding: content_encoding.clone(),
content_type: content_type.clone(),
content_length,
transfer_encoding: transfer_encoding.clone(),
response_headers: response_headers.clone(),
partial_body,
}));
}
let body = String::from_utf8_lossy(&body_bytes).into_owned();
log::debug!("SSE response length: {} bytes", body.len());
log::trace!("SSE response body: {}", &body[..body.len().min(2000)]);
let dto = match Self::parse_sse_response(&body) {
Ok(v) => v,
Err(e) => {
if let Some(t) = tracer {
t.end_span(&self.model);
}
return Err(Box::new(OpenAIClientError::Deserialize { source: e, body }));
return Err(e);
}
};
@ -179,9 +292,229 @@ impl OpenAIClient {
let result = build_llm_result(dto, tools);
if result.summary.is_empty() && result.tool_call.is_none() {
return Err(Box::new(OpenAIClientError::NoText { body }));
return Err(Box::new(OpenAIClientError::NoText {
body: format!(
"empty response from SSE stream; sse_body={}",
Self::format_body_excerpt(&body, 4000)
),
}));
}
Ok(result)
}
/// Parse an SSE response body, extracting the ResponseDTO from the `response.completed` event.
///
/// SSE format from OpenAI Responses API:
/// ```text
/// event: response.created
/// data: {"type":"response.created","response":{...}}
///
/// event: response.output_text.delta
/// data: {"type":"response.output_text.delta","delta":"Hello",...}
///
/// event: response.completed
/// data: {"type":"response.completed","response":{"output":[...],...}}
/// ```
fn parse_sse_response(body: &str) -> Result<ResponseDTO, Box<dyn Error + Send + Sync>> {
let mut current_event: Option<String> = None;
let mut completed_data: Option<String> = None;
for line in body.lines() {
if let Some(event_name) = line.strip_prefix("event:") {
current_event = Some(event_name.trim().to_string());
} else if let Some(data_payload) = line.strip_prefix("data:") {
let data = data_payload.trim_start();
if current_event.as_deref() == Some("response.completed") {
completed_data = Some(data.to_string());
break;
}
}
// Empty lines are event boundaries in SSE — reset current_event
if line.is_empty() {
current_event = None;
}
}
let data = match completed_data {
Some(d) => d,
None => {
// Fallback: maybe the body is plain JSON (non-SSE), try parsing directly
log::warn!(
"No response.completed event found in SSE stream, trying direct JSON parse"
);
return serde_json::from_str::<ResponseDTO>(body.trim()).map_err(|e| {
log::error!(
"Direct JSON parse also failed. Body prefix: {}",
&body[..body.len().min(500)]
);
Box::new(OpenAIClientError::Deserialize {
source: e,
body: format!("sse_body={}", Self::format_body_excerpt(body, 4000)),
}) as Box<dyn Error + Send + Sync>
});
}
};
// The response.completed data has shape: {"type":"response.completed","response":{...}}
let wrapper: serde_json::Value = serde_json::from_str(&data).map_err(|e| {
log::error!(
"Failed to parse response.completed JSON: {}",
&data[..data.len().min(500)]
);
Box::new(OpenAIClientError::Deserialize {
source: e,
body: format!(
"response.completed data={}; sse_body={}",
Self::format_body_excerpt(&data, 1000),
Self::format_body_excerpt(body, 4000)
),
}) as Box<dyn Error + Send + Sync>
})?;
// Extract the "response" field which contains the actual response object
let response_obj = wrapper.get("response").unwrap_or(&wrapper);
serde_json::from_value::<ResponseDTO>(response_obj.clone()).map_err(|e| {
log::error!(
"Failed to deserialize ResponseDTO from: {}",
&response_obj.to_string()[..response_obj.to_string().len().min(500)]
);
Box::new(OpenAIClientError::Deserialize {
source: e,
body: format!(
"response_obj={}; sse_body={}",
Self::format_body_excerpt(&response_obj.to_string(), 1000),
Self::format_body_excerpt(body, 4000)
),
}) as Box<dyn Error + Send + Sync>
})
}
fn format_body_excerpt(body: &str, max_len: usize) -> String {
let trimmed = body.trim();
let total = trimmed.chars().count();
if total <= max_len {
trimmed.to_string()
} else {
let excerpt: String = trimmed.chars().take(max_len).collect();
format!(
"{}...[truncated {} chars]",
excerpt,
total.saturating_sub(max_len)
)
}
}
fn format_headers_excerpt(headers: &reqwest::header::HeaderMap) -> String {
let mut rendered = String::new();
for (name, value) in headers.iter() {
if !rendered.is_empty() {
rendered.push_str(", ");
}
let value_str = value.to_str().unwrap_or("<non-utf8>");
rendered.push_str(name.as_str());
rendered.push('=');
rendered.push_str(value_str);
if rendered.len() > 2000 {
rendered.push_str("...[truncated]");
break;
}
}
if rendered.is_empty() {
"<missing>".to_string()
} else {
rendered
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_sse_text_response() {
let sse_body = "\
event: response.created\n\
data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"output\":[]}}\n\
\n\
event: response.output_text.delta\n\
data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hello\"}\n\
\n\
event: response.completed\n\
data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hello, world!\"}]}]}}\n\
\n";
let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse SSE");
let (summary, tool_call) = dto.extract_parts();
assert_eq!(summary, "Hello, world!");
assert!(tool_call.is_none());
}
#[test]
fn parse_sse_function_call_response() {
let sse_body = "\
event: response.created\n\
data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_2\",\"output\":[]}}\n\
\n\
event: response.completed\n\
data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_2\",\"output\":[{\"type\":\"function_call\",\"name\":\"shell_exec\",\"arguments\":\"{\\\"command\\\":\\\"ls\\\"}\",\"call_id\":\"call_abc\"}]}}\n\
\n";
let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse SSE");
let (summary, tool_call) = dto.extract_parts();
assert!(summary.is_empty());
let call = tool_call.expect("should have tool call");
assert_eq!(call.name, "shell_exec");
assert_eq!(call.call_id, "call_abc");
}
#[test]
fn parse_sse_mixed_text_and_function_call() {
let sse_body = "\
event: response.completed\n\
data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Let me check.\"}]},{\"type\":\"function_call\",\"name\":\"read_file\",\"arguments\":\"{}\",\"call_id\":\"call_xyz\"}]}}\n\
\n";
let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse SSE");
let (summary, tool_call) = dto.extract_parts();
assert_eq!(summary, "Let me check.");
assert!(tool_call.is_some());
}
#[test]
fn parse_sse_no_completed_event_falls_back_to_direct_json() {
// If body is plain JSON (not SSE), should still work
let json_body = "{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Direct JSON\"}]}]}";
let dto = OpenAIClient::parse_sse_response(json_body).expect("should parse direct JSON");
let (summary, _) = dto.extract_parts();
assert_eq!(summary, "Direct JSON");
}
#[test]
fn parse_sse_empty_body_errors() {
let result = OpenAIClient::parse_sse_response("");
assert!(result.is_err());
}
#[test]
fn parse_sse_ignores_earlier_events() {
// Ensure we pick up data from response.completed, not from earlier events
let sse_body = "\
event: response.created\n\
data: {\"type\":\"response.created\",\"response\":{\"id\":\"r1\",\"output\":[]}}\n\
\n\
event: response.in_progress\n\
data: {\"type\":\"response.in_progress\"}\n\
\n\
event: response.completed\n\
data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Final answer\"}]}]}}\n\
\n";
let dto = OpenAIClient::parse_sse_response(sse_body).expect("should parse");
let (summary, _) = dto.extract_parts();
assert_eq!(summary, "Final answer");
}
}

2
src/infrastructure/api_clients/openai/dto/request.rs
View file

@ -141,7 +141,7 @@ impl RequestDTO {
tool_choice: "auto".to_string(),
parallel_tool_calls: false,
store: false,
stream: false,
stream: true,
}
}
}

2
src/infrastructure/auth/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod oauth;
pub use oauth::{initiate_oauth_flow, refresh_oauth_tokens};

224
src/infrastructure/auth/oauth.rs Normal file
View file

@ -0,0 +1,224 @@
use oauth2::{
AuthUrl, AuthorizationCode, ClientId, CsrfToken, PkceCodeChallenge,
PkceCodeVerifier, RedirectUrl, RefreshToken, Scope, TokenResponse, TokenUrl,
};
use oauth2::basic::BasicClient;
use oauth2::reqwest::http_client;
use sha2::{Digest, Sha256};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
use tiny_http::{Response, Server};
const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
const AUTH_URL: &str = "https://auth.openai.com/oauth/authorize";
const TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
const REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
const SCOPES: &[&str] = &["openid", "profile", "email", "offline_access"];
pub struct OAuthTokens {
pub access_token: String,
pub refresh_token: String,
pub expires_in: i64,
pub account_id: Option<String>,
}
/// Generates a cryptographically secure random string for PKCE
fn generate_random_string(length: usize) -> String {
use rand::Rng;
const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~";
let mut rng = rand::thread_rng();
(0..length)
.map(|_| CHARSET[rng.gen_range(0..CHARSET.len())] as char)
.collect()
}
/// Initiates OAuth flow and returns tokens
pub fn initiate_oauth_flow() -> Result<OAuthTokens, String> {
// Generate PKCE code verifier and challenge
let code_verifier = generate_random_string(43);
let mut hasher = Sha256::new();
hasher.update(code_verifier.as_bytes());
use base64::Engine;
let _challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD
.encode(hasher.finalize());
let pkce_verifier = PkceCodeVerifier::new(code_verifier);
let pkce_challenge = PkceCodeChallenge::from_code_verifier_sha256(&pkce_verifier);
// Build OAuth client
let client = BasicClient::new(
ClientId::new(CLIENT_ID.to_string()),
None,
AuthUrl::new(AUTH_URL.to_string()).map_err(|e| e.to_string())?,
Some(TokenUrl::new(TOKEN_URL.to_string()).map_err(|e| e.to_string())?),
)
.set_redirect_uri(RedirectUrl::new(REDIRECT_URI.to_string()).map_err(|e| e.to_string())?);
// Generate authorization URL with OpenCode-specific parameters
let (mut auth_url, csrf_state) = client
.authorize_url(CsrfToken::new_random)
.add_scopes(SCOPES.iter().map(|s| Scope::new(s.to_string())))
.set_pkce_challenge(pkce_challenge)
.url();
// Add OpenCode-specific query parameters
auth_url.query_pairs_mut()
.append_pair("id_token_add_organizations", "true")
.append_pair("codex_cli_simplified_flow", "true")
.append_pair("originator", "drastis");
// Start local HTTP server for callback
let server = Server::http("127.0.0.1:1455")
.map_err(|e| format!("Failed to start callback server on port 1455: {}", e))?;
log::info!("Starting OAuth flow, opening browser...");
// Open browser
open::that(auth_url.to_string())
.map_err(|e| format!("Failed to open browser: {}", e))?;
// Wait for callback with timeout
let (tx, rx) = mpsc::channel();
let csrf_secret = csrf_state.secret().clone();
thread::spawn(move || {
// Listen for callback (2 minute timeout)
if let Ok(Some(request)) = server.recv_timeout(Duration::from_secs(120)) {
let url_str = format!("http://localhost{}", request.url());
// Send success HTML response
let html = r#"
<html>
<head><title>Authentication Successful</title></head>
<body>
<h1>✓ Authentication Successful</h1>
<p>You can close this window and return to the CLI.</p>
<script>setTimeout(() => window.close(), 2000);</script>
</body>
</html>
"#;
let _ = request.respond(Response::from_string(html)
.with_header(tiny_http::Header::from_bytes(&b"Content-Type"[..], &b"text/html"[..]).unwrap()));
let _ = tx.send(url_str);
}
});
// Wait for callback
let callback_url = rx.recv_timeout(Duration::from_secs(120))
.map_err(|_| "OAuth timeout: No response received within 2 minutes".to_string())?;
// Parse callback URL
let url = url::Url::parse(&callback_url)
.map_err(|e| format!("Invalid callback URL: {}", e))?;
// Extract code and state from query parameters
let mut code = None;
let mut state = None;
for (key, value) in url.query_pairs() {
match key.as_ref() {
"code" => code = Some(value.to_string()),
"state" => state = Some(value.to_string()),
_ => {}
}
}
let code = code.ok_or("No authorization code in callback")?;
let state = state.ok_or("No state in callback")?;
// Verify CSRF token
if state != csrf_secret {
return Err("CSRF token mismatch - possible security attack".to_string());
}
// Exchange code for tokens
log::info!("Exchanging authorization code for tokens...");
let token_response = client
.exchange_code(AuthorizationCode::new(code))
.set_pkce_verifier(pkce_verifier)
.request(http_client)
.map_err(|e| format!("Token exchange failed: {}", e))?;
let access_token = token_response.access_token().secret().clone();
let refresh_token = token_response.refresh_token()
.ok_or("No refresh token received")?
.secret()
.clone();
let expires_in = token_response.expires_in()
.ok_or("No expiration time received")?
.as_secs() as i64;
// Extract account ID from access token (JWT)
let account_id = extract_account_id_from_token(&access_token);
log::info!("OAuth authentication successful!");
Ok(OAuthTokens {
access_token,
refresh_token,
expires_in,
account_id,
})
}
/// Refreshes OAuth tokens using refresh token
pub fn refresh_oauth_tokens(refresh_token: &str) -> Result<OAuthTokens, String> {
let client = BasicClient::new(
ClientId::new(CLIENT_ID.to_string()),
None,
AuthUrl::new(AUTH_URL.to_string()).map_err(|e| e.to_string())?,
Some(TokenUrl::new(TOKEN_URL.to_string()).map_err(|e| e.to_string())?),
);
log::info!("Refreshing OAuth tokens...");
let token_response = client
.exchange_refresh_token(&RefreshToken::new(refresh_token.to_string()))
.request(http_client)
.map_err(|e| format!("Token refresh failed: {}", e))?;
let access_token = token_response.access_token().secret().clone();
let refresh_token = token_response.refresh_token()
.map(|t| t.secret().clone())
.unwrap_or_else(|| refresh_token.to_string());
let expires_in = token_response.expires_in()
.ok_or("No expiration time received")?
.as_secs() as i64;
let account_id = extract_account_id_from_token(&access_token);
Ok(OAuthTokens {
access_token,
refresh_token,
expires_in,
account_id,
})
}
/// Extracts ChatGPT account ID from JWT token
fn extract_account_id_from_token(token: &str) -> Option<String> {
use base64::Engine;
// JWT format: header.payload.signature
let parts: Vec<&str> = token.split('.').collect();
if parts.len() != 3 {
return None;
}
// Decode payload (base64url)
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(parts[1])
.ok()?;
// Parse JSON
let json: serde_json::Value = serde_json::from_slice(&payload).ok()?;
// Extract chatgpt_account_id
json.get("chatgpt_account_id")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
}

1
src/infrastructure/cli/actions/mod.rs
View file

@ -1,3 +1,4 @@
pub mod change_settings;
pub mod insert_openai_key;
pub mod oauth_login;
pub mod select_openai_model;

48
src/infrastructure/cli/actions/oauth_login.rs Normal file
View file

@ -0,0 +1,48 @@
use crate::infrastructure::auth::initiate_oauth_flow;
use crate::infrastructure::cli::repl::refresh_settings_from_db;
use crate::infrastructure::cli::state::UiState;
use crate::infrastructure::db::DbPool;
use crate::infrastructure::event_bus::{EventBus, UiToAgentEvent};
use crate::repository::UserSettingsRepository;
pub fn start_oauth_flow(
bus: &EventBus,
conn: &DbPool,
state: &mut UiState,
) -> Result<(), String> {
log::info!("Starting OpenAI OAuth login flow...");
// Initiate OAuth flow (blocks until complete or timeout)
let tokens = initiate_oauth_flow()?;
// Store tokens in database
let conn_guard = conn
.get()
.map_err(|e| format!("Database connection failed: {}", e))?;
let settings_repo = UserSettingsRepository::new(&*conn_guard);
settings_repo.update_oauth_tokens(
&tokens.access_token,
&tokens.refresh_token,
tokens.expires_in,
tokens.account_id.as_deref(),
)?;
log::info!("OAuth tokens stored successfully");
// Notify event controller to reload engine with OAuth credentials
let _ = bus.ui_to_agent_tx.send(UiToAgentEvent::SettingsUpdateEvent {
model: None,
openai_api_key: None,
use_behavior_trees: None,
openai_tracing_enabled: None,
web_search_enabled: None,
max_tool_calls_per_request: None,
brave_api_key: None,
});
// Refresh UI state from database
refresh_settings_from_db(conn, state)?;
Ok(())
}

52
src/infrastructure/cli/actions/select_openai_model.rs
View file

@ -35,7 +35,7 @@ pub fn handle_model_entry(
refresh_settings_from_db(conn, state)?;
state.mode = UiMode::Normal;
}
ModelEntry::OpenAi { name } => {
ModelEntry::OpenAi { name, .. } => {
update_model_selection_in_db(conn, &ModelSelection::OpenAiModel(name.clone()))?;
let _ = bus
.ui_to_agent_tx
@ -52,17 +52,7 @@ pub fn handle_model_entry(
state.mode = UiMode::Normal;
}
ModelEntry::OpenAiSelect => {
if insert_openai_key::api_key_missing(conn)? {
insert_openai_key::begin_prompt(state, false);
} else {
state.models.openai_available_status = LoadStatus::Loading;
state.openai_fetch_pending = true;
state.mode = UiMode::Popup(PopupState::OpenAiAvailable {
selected: 0,
filter: String::new(),
cursor: 0,
});
}
state.mode = UiMode::Popup(PopupState::AuthMethodSelect { selected: 0 });
}
ModelEntry::Loading | ModelEntry::Empty => {}
}
@ -70,6 +60,42 @@ pub fn handle_model_entry(
Ok(())
}
pub fn handle_auth_method_api_key(
conn: &DbPool,
state: &mut UiState,
) -> Result<(), String> {
if insert_openai_key::api_key_missing(conn)? {
insert_openai_key::begin_prompt(state, false);
} else {
state.models.openai_available_status = LoadStatus::Loading;
state.openai_fetch_pending = true;
state.mode = UiMode::Popup(PopupState::OpenAiAvailable {
selected: 0,
filter: String::new(),
cursor: 0,
});
}
Ok(())
}
pub fn handle_auth_method_oauth(
bus: &EventBus,
conn: &DbPool,
state: &mut UiState,
) -> Result<(), String> {
use crate::infrastructure::cli::actions::oauth_login;
oauth_login::start_oauth_flow(bus, conn, state)?;
// After successful OAuth, proceed to model selection
state.models.openai_available_status = LoadStatus::Loading;
state.openai_fetch_pending = true;
state.mode = UiMode::Popup(PopupState::OpenAiAvailable {
selected: 0,
filter: String::new(),
cursor: 0,
});
Ok(())
}
pub fn handle_openai_available_selection(
bus: &EventBus,
conn: &DbPool,
@ -91,4 +117,4 @@ pub fn handle_openai_available_selection(
refresh_settings_from_db(conn, state)?;
state.mode = UiMode::Normal;
Ok(())
}
}

13
src/infrastructure/cli/components/permissions.rs
View file

@ -16,7 +16,6 @@ pub fn render(
scope: &str,
selected: usize,
theme: Theme,
is_read_only: bool,
) {
let width = min(
(size.width as f32 * 0.7) as u16,
@ -58,7 +57,7 @@ pub fn render(
width: area.width,
height: 5,
};
let options = permission_options(command.is_some(), is_read_only);
let options = permission_options(command.is_some());
let items: Vec<ListItem> = options
.iter()
.map(|option| ListItem::new(Line::from(Span::raw(*option))))
@ -75,23 +74,17 @@ pub fn render(
frame.render_stateful_widget(list, options_area, &mut list_state);
}
fn permission_options(has_command: bool, is_read_only: bool) -> Vec<&'static str> {
fn permission_options(has_command: bool) -> Vec<&'static str> {
if has_command {
vec![
"[A] Allow",
"[C] Allow this Command for this Project",
"[D] Disallow",
]
} else if is_read_only {
vec![
"[A] Allow",
"[R] Allow All Reads in Project for this Session",
"[D] Disallow",
]
} else {
vec![
"[A] Allow",
"[W] Allow All Writes in Project for this Session",
"[W] Allow All Writes in this Project",
"[D] Disallow",
]
}

29
src/infrastructure/cli/components/settings_panel.rs
View file

@ -1,3 +1,4 @@
use crate::domain::AuthMethod;
use crate::infrastructure::cli::helpers::{checkbox_item, list_state, panel_block};
use crate::infrastructure::cli::theme::{Theme, PANEL_PADDING};
use ratatui::style::{Modifier, Style};
@ -13,8 +14,29 @@ pub fn render(
openai_tracing: bool,
web_search: bool,
max_tool_calls: i32,
auth_method: &AuthMethod,
oauth_token_expiry: Option<i64>,
theme: Theme,
) {
// Determine auth status display
let auth_status = match auth_method {
AuthMethod::OAuth => {
let expiry_msg = if let Some(exp) = oauth_token_expiry {
let now = chrono::Utc::now().timestamp();
if exp > now {
let mins_left = (exp - now) / 60;
format!(" (expires in {}m)", mins_left)
} else {
" (EXPIRED - re-login needed)".to_string()
}
} else {
"".to_string()
};
format!("OAuth{}", expiry_msg)
}
AuthMethod::ApiKey => "API Key".to_string(),
};
let items = vec![
checkbox_item("Behavior trees", behavior_trees),
checkbox_item("OpenAI tracing", openai_tracing),
@ -26,6 +48,13 @@ pub fn render(
Style::default().fg(theme.active),
),
])),
ListItem::new(Line::from(vec![
Span::raw("Auth: "),
Span::styled(
auth_status,
Style::default().fg(theme.active),
),
])),
];
let list = List::new(items)
.highlight_style(

97
src/infrastructure/cli/controls.rs
View file

@ -1,4 +1,4 @@
use crate::domain::permissions::PermissionDecision;
use crate::domain::permissions::UserPermissionDecision;
use crate::domain::AgentModeType;
use crate::infrastructure::cli::actions::{
change_settings, insert_openai_key, select_openai_model,
@ -24,6 +24,13 @@ pub fn handle_key_event(
state: &mut UiState,
key: KeyEvent,
) -> Result<(), String> {
if key.code == KeyCode::Char('c') && key.modifiers.contains(KeyModifiers::CONTROL) {
if state.request_in_flight.is_some() {
let _ = bus.ui_to_agent_tx.send(UiToAgentEvent::CancelRequest);
}
return Ok(());
}
let mode_snapshot = state.mode.clone();
match mode_snapshot {
UiMode::Normal => handle_normal_key(bus, conn, state, key),
@ -213,8 +220,8 @@ fn handle_normal_key(
// Enter right_container review mode for whatever is displayed
let has_file_changes = state.file_changes.is_some()
&& !state.file_changes.as_ref().unwrap().changes.is_empty();
let has_todo_list = state.todo_list.is_some()
&& !state.todo_list.as_ref().unwrap().items.is_empty();
let has_todo_list =
state.todo_list.is_some() && !state.todo_list.as_ref().unwrap().items.is_empty();
// Priority: file changes first (if both exist), otherwise whichever exists
if has_file_changes {
@ -237,8 +244,8 @@ fn handle_normal_key(
// Enter right_container review mode for whatever is displayed
let has_file_changes = state.file_changes.is_some()
&& !state.file_changes.as_ref().unwrap().changes.is_empty();
let has_todo_list = state.todo_list.is_some()
&& !state.todo_list.as_ref().unwrap().items.is_empty();
let has_todo_list =
state.todo_list.is_some() && !state.todo_list.as_ref().unwrap().items.is_empty();
// Priority: file changes first (if both exist), otherwise whichever exists
if has_file_changes {
@ -568,12 +575,17 @@ fn handle_popup_key(
}
_ => {}
},
PopupState::PermissionPrompt { selected, is_read_only, command, .. } => {
PopupState::PermissionPrompt {
selected,
is_read_only,
command,
..
} => {
let options_len = permissions::option_count();
let has_command = command.is_some();
match key.code {
KeyCode::Esc | KeyCode::Char('d') | KeyCode::Char('D') => {
submit_permission_decision(bus, popup, PermissionDecision::Ask)?;
submit_permission_decision(bus, popup, UserPermissionDecision::Deny)?;
state.mode = UiMode::Normal;
}
KeyCode::Up => {
@ -587,39 +599,45 @@ fn handle_popup_key(
}
}
KeyCode::Char('a') | KeyCode::Char('A') => {
submit_permission_decision(bus, popup, PermissionDecision::AllowOnce)?;
submit_permission_decision(bus, popup, UserPermissionDecision::AllowOnce)?;
state.mode = UiMode::Normal;
}
KeyCode::Char('c') | KeyCode::Char('C') => {
if has_command {
submit_permission_decision(bus, popup, PermissionDecision::AllowCommandForProject)?;
submit_permission_decision(
bus,
popup,
UserPermissionDecision::AlwaysAllow,
)?;
state.mode = UiMode::Normal;
}
}
KeyCode::Char('r') | KeyCode::Char('R') => {
if !has_command && *is_read_only {
submit_permission_decision(bus, popup, PermissionDecision::AllowAllReadsInSession)?;
submit_permission_decision(
bus,
popup,
UserPermissionDecision::AlwaysAllow,
)?;
state.mode = UiMode::Normal;
}
}
KeyCode::Char('w') | KeyCode::Char('W') => {
if !has_command && !*is_read_only {
submit_permission_decision(bus, popup, PermissionDecision::AllowAllWritesInSession)?;
submit_permission_decision(
bus,
popup,
UserPermissionDecision::AlwaysAllow,
)?;
state.mode = UiMode::Normal;
}
}
KeyCode::Enter => {
let decision = match *selected {
0 => PermissionDecision::AllowOnce,
1 => if has_command {
PermissionDecision::AllowCommandForProject
} else if *is_read_only {
PermissionDecision::AllowAllReadsInSession
} else {
PermissionDecision::AllowAllWritesInSession
},
2 => PermissionDecision::Ask,
_ => PermissionDecision::Ask,
0 => UserPermissionDecision::AllowOnce,
1 => UserPermissionDecision::AlwaysAllow,
2 => UserPermissionDecision::Deny,
_ => UserPermissionDecision::Deny,
};
submit_permission_decision(bus, popup, decision)?;
state.mode = UiMode::Normal;
@ -627,6 +645,29 @@ fn handle_popup_key(
_ => {}
}
}
PopupState::AuthMethodSelect { selected } => match key.code {
KeyCode::Esc => state.mode = UiMode::Popup(PopupState::ModelSelect { selected: 0 }),
KeyCode::Up => {
if *selected > 0 {
*selected -= 1;
}
}
KeyCode::Down => {
if *selected < 1 {
*selected += 1;
}
}
KeyCode::Enter => match *selected {
0 => {
select_openai_model::handle_auth_method_api_key(conn, state)?;
}
1 => {
select_openai_model::handle_auth_method_oauth(bus, conn, state)?;
}
_ => {}
},
_ => {}
},
PopupState::ContinueSelect { sessions, selected } => {
let count = sessions.len();
match key.code {
@ -644,9 +685,9 @@ fn handle_popup_key(
KeyCode::Enter => {
let session_id = sessions[*selected].id;
state.session_id = Some(session_id);
let _ = bus.ui_to_agent_tx.send(UiToAgentEvent::SessionContinueEvent {
session_id,
});
let _ = bus
.ui_to_agent_tx
.send(UiToAgentEvent::SessionContinueEvent { session_id });
state.mode = UiMode::Normal;
}
_ => {}
@ -862,7 +903,7 @@ fn handle_popup_input(state: &mut UiState, key: KeyEvent) {
fn submit_permission_decision(
bus: &EventBus,
popup: &PopupState,
decision: PermissionDecision,
decision: UserPermissionDecision,
) -> Result<(), String> {
if let PopupState::PermissionPrompt { request_id, .. } = popup {
bus.ui_to_agent_tx
@ -1328,14 +1369,12 @@ fn open_continue_popup(conn: &DbPool, state: &mut UiState) {
.and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string()))
.unwrap_or_else(|| "default".to_string());
let projects_repo =
crate::repository::ProjectsRepository::new(&*conn_guard);
let projects_repo = crate::repository::ProjectsRepository::new(&*conn_guard);
let project = projects_repo
.find_by_name(&project_name)?
.ok_or_else(|| "not_found".to_string())?;
let sessions_repo =
crate::repository::SessionsRepository::new(&*conn_guard);
let sessions_repo = crate::repository::SessionsRepository::new(&*conn_guard);
let rows = sessions_repo.find_by_project_recent(project.id, 5)?;
if rows.is_empty() {

124
src/infrastructure/cli/repl.rs
View file

@ -1,4 +1,4 @@
use crate::domain::{AgentModeType, ModelType};
use crate::domain::{AgentModeType, ModelAuthType, ModelType};
use crate::infrastructure::cli::state::{
FileChangesDisplay, LoadStatus, PopupState, ProgressEntry, ProgressKind, RequestIndicator,
RequestStatusDisplay, UiMode, UiState,
@ -180,9 +180,7 @@ fn handle_agent_event(
let can_update = state
.last_progress_update
.map(|last| {
std::time::Instant::now()
.duration_since(last)
.as_millis()
std::time::Instant::now().duration_since(last).as_millis()
>= crate::infrastructure::cli::state::MIN_PROGRESS_DISPLAY_MS
})
.unwrap_or(true); // Update immediately if no previous update
@ -236,15 +234,29 @@ fn handle_agent_event(
}
}
// Only show summary for failures, skip for success
if let Some(summary) = summary {
// Show failures via summary; for successful requests, fallback to summary
// when final_message is missing/empty (common when resuming older sessions).
if let Some(summary_text) = summary {
if status != RequestStatus::Success {
state.push_progress(ProgressEntry {
text: summary,
text: summary_text,
kind: finish_kind,
active: false,
request_id: None,
});
} else {
let final_is_empty = final_message
.as_ref()
.map(|m| m.trim().is_empty())
.unwrap_or(true);
if final_is_empty {
state.push_progress(ProgressEntry {
text: summary_text,
kind: ProgressKind::Info,
active: false,
request_id: None,
});
}
}
}
if let Some(message) = final_message {
@ -339,6 +351,7 @@ pub(crate) fn refresh_models_from_db(conn: &DbPool, state: &mut UiState) -> Resu
.map(|name| crate::infrastructure::event_bus::OpenAiModelInfo {
_id: model.id,
name,
auth_type: model.auth_type.as_str().to_string(),
})
})
.collect();
@ -362,6 +375,8 @@ pub(crate) fn refresh_settings_from_db(conn: &DbPool, state: &mut UiState) -> Re
state.settings.openai_tracing_enabled = settings.openai_tracing_enabled;
state.settings.web_search_enabled = settings.web_search_enabled;
state.settings.max_tool_calls_per_request = settings.max_tool_calls_per_request;
state.settings.auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method);
state.settings.oauth_token_expiry = settings.oauth_token_expiry;
state.settings.status = LoadStatus::Loaded;
// Get both model name and type
@ -437,10 +452,23 @@ pub(crate) fn update_model_selection_in_db(
})
});
let model_id = match existing_model {
Some(model) => model.id,
Some(model) => {
if model.auth_type != ModelAuthType::Local {
models_repo
.update_auth_type(model.id, ModelAuthType::Local)
.map_err(|e| e.to_string())?;
}
model.id
}
None => {
models_repo
.create(ModelType::Local, None, Some(&gguf_path), None)
.create(
ModelType::Local,
None,
Some(&gguf_path),
None,
ModelAuthType::Local,
)
.map_err(|e| e.to_string())?
.id
}
@ -457,14 +485,24 @@ pub(crate) fn update_model_selection_in_db(
let existing_models = models_repo
.find_by_type(ModelType::OpenAI)
.map_err(|e| e.to_string())?;
let settings = settings_repo.get_current().map_err(|e| e.to_string())?;
let auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method);
let auth_type = ModelAuthType::from_auth_method(&auth_method);
let model_id = match existing_models
.iter()
.find(|model| model.model_name.as_deref() == Some(model_name))
{
Some(model) => model.id,
Some(model) => {
if model.auth_type != auth_type {
models_repo
.update_auth_type(model.id, auth_type)
.map_err(|e| e.to_string())?;
}
model.id
}
None => {
models_repo
.create(ModelType::OpenAI, None, None, Some(model_name))
.create(ModelType::OpenAI, None, None, Some(model_name), auth_type)
.map_err(|e| e.to_string())?
.id
}
@ -491,18 +529,63 @@ pub(crate) fn fetch_openai_available_models(
.map_err(|e| format!("Failed to get connection: {}", e))?;
let settings_repo = UserSettingsRepository::new(&*conn_guard);
let settings = settings_repo.get_current().map_err(|e| e.to_string())?;
let api_key = match settings.openai_api_key {
Some(key) => key,
None => {
state.models.openai_available =
vec!["OpenAI API key is required. Please set it in settings.".to_string()];
state.models.openai_available_status = LoadStatus::Loaded;
return Ok(());
let auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method);
let auth_token = if auth_method == crate::domain::AuthMethod::OAuth {
match (&settings.oauth_access_token, &settings.oauth_account_id) {
(Some(token), account_id) => {
crate::infrastructure::api_clients::openai::client::AuthToken::OAuth {
token: token.clone(),
account_id: account_id.clone(),
}
}
_ => {
state.models.openai_available =
vec!["OAuth token not found. Please re-authenticate.".to_string()];
state.models.openai_available_status = LoadStatus::Loaded;
return Ok(());
}
}
} else {
match &settings.openai_api_key {
Some(key) if !key.trim().is_empty() => {
crate::infrastructure::api_clients::openai::client::AuthToken::ApiKey(key.clone())
}
_ => {
state.models.openai_available =
vec!["OpenAI API key is required. Please set it in settings.".to_string()];
state.models.openai_available_status = LoadStatus::Loaded;
return Ok(());
}
}
};
// OAuth tokens can't call /v1/models — return known Codex/ChatGPT models instead
if auth_method == crate::domain::AuthMethod::OAuth {
state.models.openai_available = vec![
"gpt-5.3-codex-spark".to_string(),
"gpt-5.3-codex".to_string(),
"gpt-5.2-codex".to_string(),
"gpt-5.2".to_string(),
"gpt-5.1-codex-max".to_string(),
"gpt-5.1-codex".to_string(),
"gpt-5.1".to_string(),
"gpt-5-codex".to_string(),
"gpt-5-codex-mini".to_string(),
"gpt-5".to_string(),
"gpt-4.1".to_string(),
"gpt-4.1-mini".to_string(),
"gpt-4o".to_string(),
"gpt-4o-mini".to_string(),
"o3".to_string(),
"o4-mini".to_string(),
];
state.models.openai_available_status = LoadStatus::Loaded;
return Ok(());
}
state.models.openai_available_status = LoadStatus::Loading;
match OpenAIEngine::new_general(&api_key, "gpt-4") {
match OpenAIEngine::new_general(auth_token, "gpt-4") {
Ok(openai) => match openai.fetch_available_models() {
Ok(models) => {
state.models.openai_available = models;
@ -637,6 +720,7 @@ mod tests {
state.models.openai = vec![OpenAiModelInfo {
_id: 1,
name: "gpt-4".to_string(),
auth_type: "api_key".to_string(),
}];
let entries = model_entries(&state);
assert_eq!(entries.len(), 3);
@ -734,4 +818,4 @@ mod tests {
assert_eq!(model.gguf_file_path.as_deref(), Some(canonical.as_str()));
});
}
}
}

8
src/infrastructure/cli/state.rs
View file

@ -1,4 +1,5 @@
use crate::domain::tools::FileChange;
use crate::domain::AuthMethod;
use crate::domain::AgentModeType;
use crate::infrastructure::cli::loading_bar::LoadingBar;
use crate::infrastructure::event_bus::{LocalModelInfo, OpenAiModelInfo};
@ -118,6 +119,8 @@ pub struct SettingsCache {
pub openai_tracing_enabled: bool,
pub web_search_enabled: bool,
pub max_tool_calls_per_request: i32,
pub auth_method: AuthMethod,
pub oauth_token_expiry: Option<i64>,
}
#[derive(Debug, Clone)]
@ -206,6 +209,9 @@ pub enum PopupState {
sessions: Vec<SessionPreview>,
selected: usize,
},
AuthMethodSelect {
selected: usize,
},
}
#[derive(Debug, Clone)]
@ -297,6 +303,8 @@ impl UiState {
openai_tracing_enabled: false,
web_search_enabled: false,
max_tool_calls_per_request: 50,
auth_method: AuthMethod::ApiKey,
oauth_token_expiry: None,
},
loading_bar: LoadingBar::new(),
openai_fetch_pending: false,

43
src/infrastructure/cli/views/main_view.rs
View file

@ -144,7 +144,13 @@ fn render_popup(
.map(|entry| {
let label = match entry {
ModelEntry::Local { name, .. } => format!("Local: {}", name),
ModelEntry::OpenAi { name } => format!("OpenAI: {}", name),
ModelEntry::OpenAi { name, auth_type } => {
if auth_type.trim().is_empty() {
format!("OpenAI: {}", name)
} else {
format!("OpenAI: {} ({})", name, auth_type)
}
}
ModelEntry::OpenAiSelect => "OpenAI: Select New Model".to_string(),
ModelEntry::Loading => "Loading models...".to_string(),
ModelEntry::Empty => "No models available".to_string(),
@ -272,6 +278,8 @@ fn render_popup(
*openai_tracing,
*web_search,
*max_tool_calls,
&state.settings.auth_method,
state.settings.oauth_token_expiry,
theme,
);
}
@ -423,7 +431,6 @@ fn render_popup(
paths,
scope,
selected,
is_read_only,
..
} => {
permissions::render(
@ -435,9 +442,36 @@ fn render_popup(
scope,
*selected,
theme,
*is_read_only,
);
}
PopupState::AuthMethodSelect { selected } => {
let height = 7u16;
let width = min(
(size.width as f32 * 0.7) as u16,
size.width.saturating_sub(2),
);
let area = centered_rect(size, width, height);
frame.render_widget(Clear, area);
let api_key_item = ListItem::new(Line::from(Span::styled(
"API Key",
Style::default().fg(theme.info_text),
)));
let oauth_item = ListItem::new(Line::from(Span::styled(
"ChatGPT Plus (OAuth)",
Style::default().fg(theme.info_text),
)));
let list = List::new(vec![api_key_item, oauth_item])
.highlight_style(
Style::default()
.fg(theme.active)
.add_modifier(Modifier::BOLD),
)
.block(panel_block(theme, theme.panel, PANEL_PADDING))
.highlight_symbol("> ");
let mut list_state = list_state(*selected);
frame.render_stateful_widget(list, area, &mut list_state);
}
PopupState::ContinueSelect { sessions, selected } => {
if sessions.is_empty() {
return;
@ -496,7 +530,7 @@ fn render_popup(
#[derive(Debug, Clone)]
pub enum ModelEntry {
Local { name: String, path: String },
OpenAi { name: String },
OpenAi { name: String, auth_type: String },
OpenAiSelect,
Loading,
Empty,
@ -517,6 +551,7 @@ pub fn model_entries(state: &UiState) -> Vec<ModelEntry> {
for model in &state.models.openai {
entries.push(ModelEntry::OpenAi {
name: model.name.clone(),
auth_type: model.auth_type.clone(),
});
}
entries.push(ModelEntry::OpenAiSelect);

155
src/infrastructure/db.rs
View file

@ -846,6 +846,161 @@ fn run_migrations(conn: &Connection) -> Result<(), String> {
log::info!("Database migration to version 17 completed");
}
// Migration v18: Add OAuth columns to user_settings
if current_version < 18 {
log::info!("Migrating database from version {} to 18", current_version);
// Check and add auth_method column
let auth_method_exists = conn
.prepare(
"SELECT name FROM pragma_table_info('user_settings') WHERE name = 'auth_method'",
)
.and_then(|mut stmt| {
stmt.query_map([], |row| row.get::<_, String>(0))
.map(|mut rows| rows.next().is_some())
})
.unwrap_or(false);
if !auth_method_exists {
log::info!("Adding auth_method column to user_settings table");
conn.execute(
"ALTER TABLE user_settings ADD COLUMN auth_method TEXT NOT NULL DEFAULT 'api_key'",
[],
)
.map_err(|e| format!("Failed to add auth_method column: {}", e))?;
}
// Check and add oauth_access_token column
let oauth_access_token_exists = conn
.prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_access_token'")
.and_then(|mut stmt| {
stmt.query_map([], |row| row.get::<_, String>(0))
.map(|mut rows| rows.next().is_some())
})
.unwrap_or(false);
if !oauth_access_token_exists {
log::info!("Adding oauth_access_token column to user_settings table");
conn.execute(
"ALTER TABLE user_settings ADD COLUMN oauth_access_token TEXT",
[],
)
.map_err(|e| format!("Failed to add oauth_access_token column: {}", e))?;
}
// Check and add oauth_refresh_token column
let oauth_refresh_token_exists = conn
.prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_refresh_token'")
.and_then(|mut stmt| {
stmt.query_map([], |row| row.get::<_, String>(0))
.map(|mut rows| rows.next().is_some())
})
.unwrap_or(false);
if !oauth_refresh_token_exists {
log::info!("Adding oauth_refresh_token column to user_settings table");
conn.execute(
"ALTER TABLE user_settings ADD COLUMN oauth_refresh_token TEXT",
[],
)
.map_err(|e| format!("Failed to add oauth_refresh_token column: {}", e))?;
}
// Check and add oauth_token_expiry column
let oauth_token_expiry_exists = conn
.prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_token_expiry'")
.and_then(|mut stmt| {
stmt.query_map([], |row| row.get::<_, String>(0))
.map(|mut rows| rows.next().is_some())
})
.unwrap_or(false);
if !oauth_token_expiry_exists {
log::info!("Adding oauth_token_expiry column to user_settings table");
conn.execute(
"ALTER TABLE user_settings ADD COLUMN oauth_token_expiry INTEGER",
[],
)
.map_err(|e| format!("Failed to add oauth_token_expiry column: {}", e))?;
}
// Check and add oauth_account_id column
let oauth_account_id_exists = conn
.prepare("SELECT name FROM pragma_table_info('user_settings') WHERE name = 'oauth_account_id'")
.and_then(|mut stmt| {
stmt.query_map([], |row| row.get::<_, String>(0))
.map(|mut rows| rows.next().is_some())
})
.unwrap_or(false);
if !oauth_account_id_exists {
log::info!("Adding oauth_account_id column to user_settings table");
conn.execute(
"ALTER TABLE user_settings ADD COLUMN oauth_account_id TEXT",
[],
)
.map_err(|e| format!("Failed to add oauth_account_id column: {}", e))?;
}
conn.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', '18')",
[],
)
.map_err(|e| format!("Failed to update schema_version: {}", e))?;
log::info!("Database migration to version 18 completed");
}
// Migration v19: add auth_type to models
if current_version < 19 {
log::info!("Migrating database from version {} to 19", current_version);
let auth_type_exists = conn
.prepare("SELECT name FROM pragma_table_info('models') WHERE name = 'auth_type'")
.and_then(|mut stmt| {
stmt.query_map([], |row| row.get::<_, String>(0))
.map(|mut rows| rows.next().is_some())
})
.unwrap_or(false);
if !auth_type_exists {
log::info!("Adding auth_type column to models table");
conn.execute(
"ALTER TABLE models ADD COLUMN auth_type TEXT NOT NULL DEFAULT 'local'",
[],
)
.map_err(|e| format!("Failed to add auth_type column: {}", e))?;
}
let openai_auth_type: String = conn
.query_row(
"SELECT auth_method FROM user_settings WHERE id = 1",
[],
|row| row.get(0),
)
.unwrap_or_else(|_| "api_key".to_string());
conn.execute(
"UPDATE models SET auth_type = 'local' WHERE type = 'local'",
[],
)
.map_err(|e| format!("Failed to backfill local auth_type: {}", e))?;
conn.execute(
"UPDATE models SET auth_type = ? WHERE type = 'openai' AND (auth_type IS NULL OR auth_type = '')",
params![openai_auth_type],
)
.map_err(|e| format!("Failed to backfill OpenAI auth_type: {}", e))?;
conn.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('schema_version', '19')",
[],
)
.map_err(|e| format!("Failed to update schema_version: {}", e))?;
log::info!("Database migration to version 19 completed");
}
conn.execute(
"INSERT OR IGNORE INTO user_settings (id, openai_tracing_enabled, use_behavior_trees, web_search_enabled, max_tool_calls_per_request) VALUES (1, 0, 0, 0, 50)",
[],

10
src/infrastructure/event_bus.rs
View file

@ -1,8 +1,8 @@
use crate::domain::permissions::PermissionDecision;
use crate::domain::permissions::UserPermissionDecision;
use crate::domain::todo::TodoItem;
use crate::domain::tools::FileChange;
use crate::domain::AgentModeType;
use crossbeam_channel::{unbounded, Receiver, Sender};
use crate::domain::todo::TodoItem;
#[derive(Debug, Clone)]
pub enum StepPhase {
@ -33,6 +33,7 @@ pub struct LocalModelInfo {
pub struct OpenAiModelInfo {
pub _id: i64,
pub name: String,
pub auth_type: String,
}
#[derive(Debug, Clone)]
@ -48,12 +49,13 @@ pub enum UiToAgentEvent {
mode: AgentModeType, // NEW: Pass mode to agent
session_id: Option<i64>,
},
CancelRequest,
SessionContinueEvent {
session_id: i64,
},
PermissionUpdateEvent {
request_id: u64,
decision: PermissionDecision,
decision: UserPermissionDecision,
},
ShutdownEvent,
SettingsUpdateEvent {
@ -109,7 +111,7 @@ pub enum AgentToUiEvent {
#[derive(Debug, Clone)]
pub struct PermissionUpdate {
pub request_id: u64,
pub decision: PermissionDecision,
pub decision: UserPermissionDecision,
}
#[derive(Clone, Debug)]

172
src/infrastructure/event_controller.rs
View file

@ -1,13 +1,12 @@
use crate::domain::permissions::{PermissionConfig, PermissionDecision, PermissionPrompter};
use crate::domain::permissions::{PermissionConfig, PermissionPrompter, UserPermissionDecision};
use crate::domain::{CancellationToken, SessionService, StartupService};
use crate::infrastructure::db::DbPool;
use crate::infrastructure::event_bus::{
AgentToUiEvent, EventBus, ImageAttachment, PermissionUpdate, RequestStatus, StepPhase, UiToAgentEvent,
AgentToUiEvent, EventBus, ImageAttachment, PermissionUpdate, RequestStatus, StepPhase,
UiToAgentEvent,
};
use crate::infrastructure::inference::InferenceEngine;
use crate::infrastructure::init::{
apply_model_selection, update_openai_api_key,
};
use crate::infrastructure::init::{apply_model_selection, update_openai_api_key};
use crate::repository::UserSettingsRepository;
use crate::{domain, infrastructure};
use crossbeam_channel::{select, unbounded, Receiver, Sender};
@ -18,11 +17,7 @@ use std::sync::Arc;
use std::thread;
/// Helper function to send a failure event to the UI
fn send_failure_event(
tx: &Sender<AgentToUiEvent>,
request_id: u64,
error: impl Into<String>,
) {
fn send_failure_event(tx: &Sender<AgentToUiEvent>, request_id: u64, error: impl Into<String>) {
let _ = tx.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Failure,
@ -85,7 +80,12 @@ pub struct EventController {
cancel_token: CancellationToken,
permission_response_tx: Option<Sender<PermissionUpdate>>,
request_counter: u64,
pending_request: Option<(String, Vec<ImageAttachment>, crate::domain::AgentModeType, Option<i64>)>,
pending_request: Option<(
String,
Vec<ImageAttachment>,
crate::domain::AgentModeType,
Option<i64>,
)>,
}
impl EventController {
@ -188,20 +188,74 @@ impl EventController {
}
};
if let Some(last_request) = requests.last() {
let label = request_label(&last_request.prompt);
let status = match &last_request.result_summary {
let steps_repo = crate::repository::SessionRequestStepsRepository::new(self.conn.clone());
for request in &requests {
let label = request_label(&request.prompt);
let status = match &request.result_summary {
Some(summary) if summary.starts_with("Success") => RequestStatus::Success,
Some(_) => RequestStatus::Failure,
None => RequestStatus::Failure,
};
let summary = last_request.result_summary.clone();
let summary = request.result_summary.clone();
let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::RequestStartedEvent {
request_id: 0,
label,
prompt: last_request.prompt.clone(),
prompt: request.prompt.clone(),
});
let steps = match steps_repo.load_steps_for_request(request.id) {
Ok(steps) => steps,
Err(e) => {
log::warn!(
"Failed to load steps for request {}: {}",
request.id,
e
);
Vec::new()
}
};
if steps.is_empty() {
if let Some(log) = &request._steps_log {
for line in log.lines() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent {
step_name: "assistant_response".to_string(),
phase: StepPhase::Before,
summary: trimmed.to_string(),
});
let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent {
step_name: "assistant_response".to_string(),
phase: StepPhase::After,
summary: trimmed.to_string(),
});
}
}
} else {
for step in steps {
let step_name = step
.tool_name
.clone()
.unwrap_or_else(|| step.step_type.clone());
let summary = step.summary.clone();
let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent {
step_name: step_name.clone(),
phase: StepPhase::Before,
summary: summary.clone(),
});
let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::ProgressEvent {
step_name,
phase: StepPhase::After,
summary,
});
}
}
let _ = self.bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent {
request_id: 0,
status,
@ -299,6 +353,12 @@ impl EventController {
let _ = worker_status_tx.send(());
});
}
UiToAgentEvent::CancelRequest => {
if worker_running {
log::info!("Cancelling current request via Ctrl+C");
self.cancel_token.cancel();
}
}
UiToAgentEvent::PermissionUpdateEvent {
request_id,
decision,
@ -332,7 +392,8 @@ impl EventController {
match apply_model_selection(&self.conn, selection) {
Ok(engine) => {
self.engine = Some(engine);
self.current_session_id = None;
// Keep current session when switching models so ongoing
// conversation continues with a different model.
}
Err(err) => {
send_failure_and_continue!(self, err.to_string());
@ -443,7 +504,7 @@ fn run_request_worker(
use_behavior_trees,
PermissionConfig::default(),
prompter,
event_bus.agent_to_ui_tx.clone()
event_bus.agent_to_ui_tx.clone(),
)
.map_err(|e| format!("Failed to create session service: {}", e))?;
@ -462,32 +523,58 @@ fn run_request_worker(
} else {
chain.fail_reason
};
let _ = event_bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Failure,
summary: Some(reason),
final_message: None,
});
let _ = event_bus
.agent_to_ui_tx
.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Failure,
summary: Some(reason),
final_message: None,
});
} else {
let _ = event_bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Success,
summary: Some(chain.get_summary()),
final_message: chain.final_message.map(|msg| msg.to_string()),
});
let _ = event_bus
.agent_to_ui_tx
.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Success,
summary: Some(chain.get_summary()),
final_message: chain.final_message.map(|msg| msg.to_string()),
});
}
Ok(())
}
Err(domain::session::ServiceError::Workflow(domain::workflow::Error::Cancelled)) => {
let _ = event_bus.agent_to_ui_tx.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Cancelled,
summary: Some("Cancelled".to_string()),
final_message: None,
});
let _ = event_bus
.agent_to_ui_tx
.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Cancelled,
summary: Some("Cancelled".to_string()),
final_message: None,
});
Ok(())
}
Err(err) => Err(format!("Workflow error: {}", err)),
Err(err) => {
let err_str = err.to_string();
// Check if error is due to expired OAuth token
if err_str.contains("AuthExpired") || err_str.contains("OAuth token expired") {
let _ = event_bus
.agent_to_ui_tx
.send(AgentToUiEvent::RequestFinishedEvent {
request_id,
status: RequestStatus::Failure,
summary: Some(
"OAuth token expired. Please run /oauth-login to re-authenticate."
.to_string(),
),
final_message: None,
});
return Ok(());
}
Err(format!("Workflow error: {}", err))
}
}
}
@ -533,7 +620,7 @@ fn emit_chain_progress(
#[cfg(test)]
mod tests {
use super::*;
use crate::domain::permissions::{PermissionDecision, PermissionRequest, PermissionScope};
use crate::domain::permissions::{PermissionRequest, PermissionScope, UserPermissionDecision};
use crate::infrastructure::db;
use crate::infrastructure::event_bus::{ModelSelection, UiToAgentEvent};
use crate::repository::{MetaRepository, ModelsRepository, UserSettingsRepository};
@ -646,6 +733,7 @@ mod tests {
assert!(model.is_some());
let model = model.unwrap();
assert!(model._api_key.is_none());
assert_eq!(model.auth_type, domain::ModelAuthType::ApiKey);
let settings_repo = UserSettingsRepository::new(&*conn_guard);
let settings = settings_repo.get_current().expect("settings");
@ -663,7 +751,7 @@ mod tests {
vec![],
PermissionScope::Project,
Some(1),
false, // is_read_only
false, // is_read_only
PathBuf::from("/tmp"),
);
@ -681,11 +769,11 @@ mod tests {
let _ = update_tx.send(PermissionUpdate {
request_id,
decision: PermissionDecision::AllowOnce,
decision: UserPermissionDecision::AllowOnce,
});
let decision = handler.join().expect("join handler");
assert_eq!(decision, PermissionDecision::AllowOnce);
assert_eq!(decision, UserPermissionDecision::AllowOnce);
}
}
@ -709,7 +797,7 @@ impl PermissionPrompter for BusPermissionPrompter {
fn ask_permission(
&self,
request: &domain::permissions::PermissionRequest,
) -> Result<PermissionDecision, crate::utils::AskError> {
) -> Result<UserPermissionDecision, crate::utils::AskError> {
let request_id = self.counter.fetch_add(1, Ordering::SeqCst);
let scope = match request.scope {
domain::permissions::PermissionScope::Session => "session",

29
src/infrastructure/inference/openai.rs
View file

@ -1,5 +1,6 @@
use super::{InferenceEngine, LLMInferenceResult};
use crate::domain::ModelType;
use crate::infrastructure::api_clients::openai::client::AuthToken;
use crate::infrastructure::api_clients::openai::OpenAIClient;
use crate::infrastructure::InfaError;
use openai_api_rust::models::ModelsApi;
@ -12,11 +13,17 @@ pub struct OpenAIEngine {
}
impl OpenAIEngine {
/// Create a new OpenAI engine with the given API key and model
pub fn new(api_key: &str, model: &str) -> Result<Arc<dyn InferenceEngine>, String> {
let auth = Auth::new(api_key);
/// Create a new OpenAI engine with the given auth token and model
pub fn new(auth_token: AuthToken, model: &str) -> Result<Arc<dyn InferenceEngine>, String> {
// Extract API key string for openai_api_rust client (used for model listing)
let api_key_str = match &auth_token {
AuthToken::ApiKey(key) => key.clone(),
AuthToken::OAuth { token, .. } => token.clone(),
};
let auth = Auth::new(&api_key_str);
let client = OpenAI::new(auth, "https://api.openai.com/v1/");
let responses_client = OpenAIClient::new(api_key.to_string(), model.to_string());
let responses_client = OpenAIClient::new(auth_token, model.to_string());
let engine = Arc::new(Self {
client,
@ -26,11 +33,17 @@ impl OpenAIEngine {
Ok(engine as Arc<dyn InferenceEngine>)
}
/// Create a new OpenAI engine with the given API key and model
pub fn new_general(api_key: &str, model: &str) -> Result<Self, String> {
let auth = Auth::new(api_key);
/// Create a new OpenAI engine with the given auth token and model
pub fn new_general(auth_token: AuthToken, model: &str) -> Result<Self, String> {
// Extract API key string for openai_api_rust client (used for model listing)
let api_key_str = match &auth_token {
AuthToken::ApiKey(key) => key.clone(),
AuthToken::OAuth { token, .. } => token.clone(),
};
let auth = Auth::new(&api_key_str);
let client = OpenAI::new(auth, "https://api.openai.com/v1/");
let responses_client = OpenAIClient::new(api_key.to_string(), model.to_string());
let responses_client = OpenAIClient::new(auth_token, model.to_string());
let engine = Self {
client,

152
src/infrastructure/init.rs
View file

@ -1,4 +1,6 @@
use crate::domain::ModelType;
use crate::domain::{ModelAuthType, ModelType};
use crate::infrastructure::api_clients::openai::client::AuthToken;
use crate::infrastructure::auth::refresh_oauth_tokens;
use crate::infrastructure::db;
use crate::infrastructure::db::DbPool;
use crate::infrastructure::event_bus::{LocalModelInfo, ModelSelection};
@ -30,6 +32,12 @@ pub enum InitError {
#[error("OpenAI model missing api_key")]
MissingApiKey,
#[error("OAuth token missing")]
MissingOAuthToken,
#[error("OAuth token expired - please re-authenticate")]
OAuthExpired,
#[error("failed to load model engine: {0}")]
ModelLoadError(String),
@ -152,11 +160,63 @@ impl InfrastructureInitializer {
ModelType::OpenAI => {
let settings_repo = UserSettingsRepository::new(conn);
let settings = settings_repo.get_current().map_err(InitError::Repository)?;
let api_key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?;
// Use saved model_name, fallback to gpt-4 for backward compatibility
let model_name = model.model_name.as_deref().unwrap_or("gpt-4");
use crate::infrastructure::inference::openai::OpenAIEngine;
OpenAIEngine::new(&api_key, model_name).map_err(|e| InitError::ModelLoadError(e))
let auth_token = match model.auth_type {
ModelAuthType::OAuth => {
let mut access_token = settings
.oauth_access_token
.ok_or(InitError::MissingOAuthToken)?;
let refresh_token_val = settings
.oauth_refresh_token
.ok_or(InitError::MissingOAuthToken)?;
let account_id = settings.oauth_account_id;
// Check if token is expired and refresh if needed
if settings
.oauth_token_expiry
.map(|exp| chrono::Utc::now().timestamp() >= exp)
.unwrap_or(true)
{
log::info!("OAuth token expired at startup, refreshing...");
let new_tokens =
refresh_oauth_tokens(&refresh_token_val).map_err(|e| {
InitError::ModelLoadError(format!(
"Token refresh failed: {}",
e
))
})?;
settings_repo
.update_oauth_tokens(
&new_tokens.access_token,
&new_tokens.refresh_token,
new_tokens.expires_in,
new_tokens.account_id.as_deref(),
)
.map_err(InitError::Repository)?;
access_token = new_tokens.access_token;
}
AuthToken::OAuth {
token: access_token,
account_id,
}
}
ModelAuthType::ApiKey => {
let api_key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?;
AuthToken::ApiKey(api_key)
}
ModelAuthType::Local => {
return Err(InitError::ModelLoadError(
"OpenAI model has local auth_type".to_string(),
));
}
};
OpenAIEngine::new(auth_token, model_name).map_err(|e| InitError::ModelLoadError(e))
}
}
}
@ -290,6 +350,11 @@ pub fn apply_model_selection(
.unwrap_or_else(|_| PathBuf::from(existing_path));
if existing_abs == canonical_path {
model_id = Some(model.id);
if model.auth_type != ModelAuthType::Local {
models_repo
.update_auth_type(model.id, ModelAuthType::Local)
.map_err(InitError::Repository)?;
}
break;
}
}
@ -299,7 +364,13 @@ pub fn apply_model_selection(
Some(id) => id,
None => {
let model = models_repo
.create(ModelType::Local, None, Some(&gguf_path), None)
.create(
ModelType::Local,
None,
Some(&gguf_path),
None,
ModelAuthType::Local,
)
.map_err(InitError::Repository)?;
model.id
}
@ -324,12 +395,70 @@ pub fn apply_model_selection(
.map_err(InitError::Repository)?;
let settings = settings_repo.get_current().map_err(InitError::Repository)?;
let resolved_api_key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?;
let auth_method = crate::domain::AuthMethod::from_str(&settings.auth_method);
let auth_type = ModelAuthType::from_auth_method(&auth_method);
// Determine auth token based on auth method
let auth_token = match auth_type {
ModelAuthType::OAuth => {
let mut access_token = settings
.oauth_access_token
.ok_or(InitError::MissingOAuthToken)?;
let refresh_token_val = settings
.oauth_refresh_token
.ok_or(InitError::MissingOAuthToken)?;
let account_id = settings.oauth_account_id;
// Check if token is expired and refresh if needed
if settings
.oauth_token_expiry
.map(|exp| chrono::Utc::now().timestamp() >= exp)
.unwrap_or(true)
{
log::info!("OAuth token expired, refreshing...");
let new_tokens = refresh_oauth_tokens(&refresh_token_val).map_err(|e| {
InitError::ModelLoadError(format!("Token refresh failed: {}", e))
})?;
// Update database with new tokens
settings_repo
.update_oauth_tokens(
&new_tokens.access_token,
&new_tokens.refresh_token,
new_tokens.expires_in,
new_tokens.account_id.as_deref(),
)
.map_err(InitError::Repository)?;
access_token = new_tokens.access_token;
}
AuthToken::OAuth {
token: access_token,
account_id: account_id,
}
}
ModelAuthType::ApiKey => {
let key = settings.openai_api_key.ok_or(InitError::MissingApiKey)?;
AuthToken::ApiKey(key)
}
ModelAuthType::Local => {
return Err(InitError::ModelLoadError(
"OpenAI model has local auth_type".to_string(),
));
}
};
let mut model_id = None;
for model in &existing_models {
if model.model_name.as_deref() == Some(&model_name) {
model_id = Some(model.id);
if model.auth_type != auth_type {
models_repo
.update_auth_type(model.id, auth_type)
.map_err(InitError::Repository)?;
}
break;
}
}
@ -342,6 +471,11 @@ pub fn apply_model_selection(
models_repo
.update_model_name(existing.id, Some(&model_name))
.map_err(InitError::Repository)?;
if existing.auth_type != auth_type {
models_repo
.update_auth_type(existing.id, auth_type)
.map_err(InitError::Repository)?;
}
model_id = Some(existing.id);
}
}
@ -350,14 +484,14 @@ pub fn apply_model_selection(
Some(id) => id,
None => {
models_repo
.create(ModelType::OpenAI, None, None, Some(&model_name))
.create(ModelType::OpenAI, None, None, Some(&model_name), auth_type)
.map_err(InitError::Repository)?
.id
}
};
let engine = OpenAIEngine::new(&resolved_api_key, &model_name)
.map_err(InitError::ModelLoadError)?;
let engine =
OpenAIEngine::new(auth_token, &model_name).map_err(InitError::ModelLoadError)?;
meta_repo
.set_last_used_model_id(model_id)

1
src/infrastructure/mod.rs
View file

@ -1,4 +1,5 @@
mod api_clients;
pub mod auth;
pub mod cli;
pub mod db;
pub mod event_bus;

46
src/repository/models.rs
View file

@ -1,4 +1,4 @@
use crate::domain::ModelType;
use crate::domain::{ModelAuthType, ModelType};
use rusqlite::{params, Connection, Row};
/// Raw database row for models table
@ -9,6 +9,7 @@ pub struct ModelRow {
pub _api_key: Option<String>,
pub gguf_file_path: Option<String>,
pub model_name: Option<String>,
pub auth_type: ModelAuthType,
pub _date_added: String,
}
@ -16,6 +17,17 @@ impl ModelRow {
fn from_row(row: &Row) -> Result<Self, rusqlite::Error> {
let model_type_str: String = row.get(1)?;
let model_type = ModelType::from_str(&model_type_str).unwrap_or(ModelType::Local);
let auth_type_str: Option<String> = row.get(5)?;
let auth_type = match auth_type_str.as_deref() {
Some(value) => ModelAuthType::from_str(value),
None => {
if model_type == ModelType::Local {
ModelAuthType::Local
} else {
ModelAuthType::ApiKey
}
}
};
Ok(Self {
id: row.get(0)?,
@ -23,7 +35,8 @@ impl ModelRow {
_api_key: row.get(2)?,
gguf_file_path: row.get(3)?,
model_name: row.get(4)?,
_date_added: row.get(5)?,
auth_type,
_date_added: row.get(6)?,
})
}
}
@ -40,7 +53,7 @@ impl<'a> ModelsRepository<'a> {
pub fn find_by_id(&self, id: i64) -> Result<Option<ModelRow>, String> {
let mut stmt = self
.conn
.prepare("SELECT id, type, api_key, gguf_file_path, model_name, date_added FROM models WHERE id = ?")
.prepare("SELECT id, type, api_key, gguf_file_path, model_name, auth_type, date_added FROM models WHERE id = ?")
.map_err(|e| e.to_string())?;
let result = stmt
@ -54,7 +67,7 @@ impl<'a> ModelsRepository<'a> {
pub fn find_by_type(&self, model_type: ModelType) -> Result<Vec<ModelRow>, String> {
let mut stmt = self
.conn
.prepare("SELECT id, type, api_key, gguf_file_path, model_name, date_added FROM models WHERE type = ? ORDER BY date_added DESC")
.prepare("SELECT id, type, api_key, gguf_file_path, model_name, auth_type, date_added FROM models WHERE type = ? ORDER BY date_added DESC")
.map_err(|e| e.to_string())?;
let rows = stmt
@ -75,13 +88,21 @@ impl<'a> ModelsRepository<'a> {
api_key: Option<&str>,
gguf_file_path: Option<&str>,
model_name: Option<&str>,
auth_type: ModelAuthType,
) -> Result<ModelRow, String> {
let date_added = chrono_now();
self.conn
.execute(
"INSERT INTO models (type, api_key, gguf_file_path, model_name, date_added) VALUES (?, ?, ?, ?, ?)",
params![model_type.as_str(), api_key, gguf_file_path, model_name, date_added],
"INSERT INTO models (type, api_key, gguf_file_path, model_name, auth_type, date_added) VALUES (?, ?, ?, ?, ?, ?)",
params![
model_type.as_str(),
api_key,
gguf_file_path,
model_name,
auth_type.as_str(),
date_added
],
)
.map_err(|e| e.to_string())?;
@ -93,10 +114,23 @@ impl<'a> ModelsRepository<'a> {
_api_key: api_key.map(|s| s.to_string()),
gguf_file_path: gguf_file_path.map(|s| s.to_string()),
model_name: model_name.map(|s| s.to_string()),
auth_type,
_date_added: date_added,
})
}
pub fn update_auth_type(&self, id: i64, auth_type: ModelAuthType) -> Result<bool, String> {
let rows_affected = self
.conn
.execute(
"UPDATE models SET auth_type = ? WHERE id = ?",
params![auth_type.as_str(), id],
)
.map_err(|e| e.to_string())?;
Ok(rows_affected > 0)
}
pub fn update_model_name(&self, id: i64, model_name: Option<&str>) -> Result<bool, String> {
let rows_affected = self
.conn

59
src/repository/user_settings.rs
View file

@ -10,6 +10,11 @@ pub struct UserSettingsRow {
pub web_search_enabled: bool,
pub brave_api_key: Option<String>,
pub max_tool_calls_per_request: i32,
pub auth_method: String,
pub oauth_access_token: Option<String>,
pub oauth_refresh_token: Option<String>,
pub oauth_token_expiry: Option<i64>,
pub oauth_account_id: Option<String>,
}
impl UserSettingsRow {
@ -23,6 +28,13 @@ impl UserSettingsRow {
web_search_enabled: row.get::<_, i64>(5)? != 0,
brave_api_key: row.get(6)?,
max_tool_calls_per_request: row.get(7)?,
auth_method: row
.get::<_, Option<String>>(8)?
.unwrap_or_else(|| "api_key".to_string()),
oauth_access_token: row.get(9)?,
oauth_refresh_token: row.get(10)?,
oauth_token_expiry: row.get(11)?,
oauth_account_id: row.get(12)?,
})
}
}
@ -46,7 +58,7 @@ impl<'a> UserSettingsRepository<'a> {
let mut stmt = self
.conn
.prepare(
"SELECT id, openai_api_key, openai_tracing_enabled, use_behavior_trees, current_model_id, web_search_enabled, brave_api_key, max_tool_calls_per_request FROM user_settings WHERE id = ?",
"SELECT id, openai_api_key, openai_tracing_enabled, use_behavior_trees, current_model_id, web_search_enabled, brave_api_key, max_tool_calls_per_request, auth_method, oauth_access_token, oauth_refresh_token, oauth_token_expiry, oauth_account_id FROM user_settings WHERE id = ?",
)
.map_err(|e| e.to_string())?;
@ -62,7 +74,7 @@ impl<'a> UserSettingsRepository<'a> {
self.ensure_row()?;
self.conn
.execute(
"UPDATE user_settings SET openai_api_key = ? WHERE id = 1",
"UPDATE user_settings SET openai_api_key = ?, auth_method = 'api_key' WHERE id = 1",
params![api_key],
)
.map_err(|e| e.to_string())?;
@ -114,6 +126,49 @@ impl<'a> UserSettingsRepository<'a> {
Ok(())
}
pub fn update_oauth_tokens(
&self,
access_token: &str,
refresh_token: &str,
expires_in: i64,
account_id: Option<&str>,
) -> Result<(), String> {
self.ensure_row()?;
let expiry = chrono::Utc::now().timestamp() + expires_in;
self.conn
.execute(
"UPDATE user_settings SET
oauth_access_token = ?,
oauth_refresh_token = ?,
oauth_token_expiry = ?,
oauth_account_id = ?,
auth_method = 'oauth',
openai_api_key = NULL
WHERE id = 1",
params![access_token, refresh_token, expiry, account_id],
)
.map_err(|e| e.to_string())?;
Ok(())
}
#[allow(dead_code)]
pub fn clear_oauth_tokens(&self) -> Result<(), String> {
self.ensure_row()?;
self.conn
.execute(
"UPDATE user_settings SET
oauth_access_token = NULL,
oauth_refresh_token = NULL,
oauth_token_expiry = NULL,
oauth_account_id = NULL,
auth_method = 'api_key'
WHERE id = 1",
[],
)
.map_err(|e| e.to_string())?;
Ok(())
}
fn update_bool(&self, column: &str, enabled: bool) -> Result<(), String> {
self.ensure_row()?;
let sql = format!("UPDATE user_settings SET {} = ? WHERE id = 1", column);