Compare commits
No commits in common. "d90052cfd1dbfee68127a2dbafcb79d2dc5e6841" and "37f244e46b58149768ec3c4b6de2d0943a658af5" have entirely different histories.
d90052cfd1
...
37f244e46b
8 changed files with 135 additions and 324 deletions
13
Cargo.lock
generated
13
Cargo.lock
generated
|
|
@ -217,7 +217,6 @@ dependencies = [
|
||||||
"axum",
|
"axum",
|
||||||
"axum-core",
|
"axum-core",
|
||||||
"bytes",
|
"bytes",
|
||||||
"cookie",
|
|
||||||
"fastrand",
|
"fastrand",
|
||||||
"futures-util",
|
"futures-util",
|
||||||
"headers",
|
"headers",
|
||||||
|
|
@ -326,7 +325,6 @@ dependencies = [
|
||||||
"axum-extra",
|
"axum-extra",
|
||||||
"base64",
|
"base64",
|
||||||
"chrono",
|
"chrono",
|
||||||
"cookie",
|
|
||||||
"dashmap",
|
"dashmap",
|
||||||
"jsonwebtoken",
|
"jsonwebtoken",
|
||||||
"migration",
|
"migration",
|
||||||
|
|
@ -488,17 +486,6 @@ version = "0.9.6"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
|
checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "cookie"
|
|
||||||
version = "0.18.1"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747"
|
|
||||||
dependencies = [
|
|
||||||
"percent-encoding",
|
|
||||||
"time",
|
|
||||||
"version_check",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "core-foundation"
|
name = "core-foundation"
|
||||||
version = "0.9.4"
|
version = "0.9.4"
|
||||||
|
|
|
||||||
|
|
@ -13,7 +13,7 @@ reqwest = { version = "0.12", features = ["json"] }
|
||||||
serde = { version = "1.0", features = ["derive"] }
|
serde = { version = "1.0", features = ["derive"] }
|
||||||
serde_json = "1.0"
|
serde_json = "1.0"
|
||||||
axum = "0.7"
|
axum = "0.7"
|
||||||
tower-http = { version = "0.5", features = ["cors", "set-header", "limit"] }
|
tower-http = { version = "0.5", features = ["cors"] }
|
||||||
sea-orm = { version = "1.1", features = ["sqlx-postgres", "runtime-tokio-rustls", "macros"] }
|
sea-orm = { version = "1.1", features = ["sqlx-postgres", "runtime-tokio-rustls", "macros"] }
|
||||||
sea-orm-migration = "1.1"
|
sea-orm-migration = "1.1"
|
||||||
uuid = { version = "1.8", features = ["v4", "serde"] }
|
uuid = { version = "1.8", features = ["v4", "serde"] }
|
||||||
|
|
@ -22,5 +22,4 @@ tokio-cron-scheduler = "0.15.1"
|
||||||
dashmap = "6.1.0"
|
dashmap = "6.1.0"
|
||||||
jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] }
|
jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] }
|
||||||
base64 = "0.22.1"
|
base64 = "0.22.1"
|
||||||
axum-extra = { version = "0.9", features = ["typed-header", "cookie"] }
|
axum-extra = { version = "0.9", features = ["typed-header"] }
|
||||||
cookie = "0.18"
|
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,8 @@ const state = {
|
||||||
selectedRunId: null,
|
selectedRunId: null,
|
||||||
currentView: 'dashboard', // 'dashboard' or 'task'
|
currentView: 'dashboard', // 'dashboard' or 'task'
|
||||||
isEditing: false,
|
isEditing: false,
|
||||||
isAuthenticated: false
|
token: localStorage.getItem('auth_token'),
|
||||||
|
refreshToken: localStorage.getItem('refresh_token')
|
||||||
};
|
};
|
||||||
|
|
||||||
// DOM elements
|
// DOM elements
|
||||||
|
|
@ -57,15 +58,29 @@ const presetBtns = document.querySelectorAll('.btn-preset');
|
||||||
|
|
||||||
// Wrapper for fetch to include Authorization header
|
// Wrapper for fetch to include Authorization header
|
||||||
async function fetchWithAuth(url, options = {}) {
|
async function fetchWithAuth(url, options = {}) {
|
||||||
let response = await fetch(url, { ...options, credentials: 'include' });
|
if (!state.token) {
|
||||||
|
showLogin();
|
||||||
|
throw new Error('Not authenticated');
|
||||||
|
}
|
||||||
|
|
||||||
if (response.status === 401) {
|
const headers = {
|
||||||
|
...options.headers,
|
||||||
|
'Authorization': `Bearer ${state.token}`
|
||||||
|
};
|
||||||
|
|
||||||
|
let response = await fetch(url, { ...options, headers });
|
||||||
|
|
||||||
|
if (response.status === 401 && state.refreshToken) {
|
||||||
// Try to refresh token
|
// Try to refresh token
|
||||||
try {
|
try {
|
||||||
const success = await attemptTokenRefresh();
|
const success = await attemptTokenRefresh();
|
||||||
if (success) {
|
if (success) {
|
||||||
// Retry original request
|
// Retry original request with new token
|
||||||
response = await fetch(url, { ...options, credentials: 'include' });
|
const newHeaders = {
|
||||||
|
...options.headers,
|
||||||
|
'Authorization': `Bearer ${state.token}`
|
||||||
|
};
|
||||||
|
response = await fetch(url, { ...options, headers: newHeaders });
|
||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error('Token refresh failed:', error);
|
console.error('Token refresh failed:', error);
|
||||||
|
|
@ -85,12 +100,20 @@ async function attemptTokenRefresh() {
|
||||||
const response = await fetch(`${API_URL}/auth/refresh`, {
|
const response = await fetch(`${API_URL}/auth/refresh`, {
|
||||||
method: 'POST',
|
method: 'POST',
|
||||||
headers: { 'Content-Type': 'application/json' },
|
headers: { 'Content-Type': 'application/json' },
|
||||||
credentials: 'include',
|
body: JSON.stringify({ refresh_token: state.refreshToken })
|
||||||
body: JSON.stringify({ refresh_token: '' })
|
|
||||||
});
|
});
|
||||||
|
|
||||||
if (response.ok) {
|
if (response.ok) {
|
||||||
return true;
|
const data = await response.json();
|
||||||
|
if (data.access_token) {
|
||||||
|
state.token = data.access_token;
|
||||||
|
localStorage.setItem('auth_token', data.access_token);
|
||||||
|
if (data.refresh_token) {
|
||||||
|
state.refreshToken = data.refresh_token;
|
||||||
|
localStorage.setItem('refresh_token', data.refresh_token);
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error('Error during token refresh:', error);
|
console.error('Error during token refresh:', error);
|
||||||
|
|
@ -130,7 +153,7 @@ async function fetchTasks() {
|
||||||
const task = state.tasks.find((t) => t.id === state.selectedTaskId);
|
const task = state.tasks.find((t) => t.id === state.selectedTaskId);
|
||||||
if (task) {
|
if (task) {
|
||||||
if (shouldFollowLatest && task.runs && task.runs.length > 0) {
|
if (shouldFollowLatest && task.runs && task.runs.length > 0) {
|
||||||
state.selectedRunId = task.runs[0].id;
|
state.selectedRunId = task.runs[task.runs.length - 1].id;
|
||||||
}
|
}
|
||||||
|
|
||||||
renderRunHistory(task);
|
renderRunHistory(task);
|
||||||
|
|
@ -154,7 +177,13 @@ async function fetchRecentRuns() {
|
||||||
|
|
||||||
function renderTaskList() {
|
function renderTaskList() {
|
||||||
const sortedTasks = [...state.tasks].sort((a, b) => {
|
const sortedTasks = [...state.tasks].sort((a, b) => {
|
||||||
return new Date(b.created_at) - new Date(a.created_at);
|
const aDate = a.runs && a.runs.length > 0
|
||||||
|
? new Date(a.runs[0].created_at)
|
||||||
|
: new Date(a.created_at);
|
||||||
|
const bDate = b.runs && b.runs.length > 0
|
||||||
|
? new Date(b.runs[0].created_at)
|
||||||
|
: new Date(b.created_at);
|
||||||
|
return bDate - aDate;
|
||||||
});
|
});
|
||||||
|
|
||||||
taskListEl.innerHTML = sortedTasks
|
taskListEl.innerHTML = sortedTasks
|
||||||
|
|
@ -194,8 +223,8 @@ function selectTask(id, runId = null) {
|
||||||
if (runId) {
|
if (runId) {
|
||||||
state.selectedRunId = runId;
|
state.selectedRunId = runId;
|
||||||
} else if (task.runs && task.runs.length > 0) {
|
} else if (task.runs && task.runs.length > 0) {
|
||||||
// Default to latest run if not specified (index 0 is newest)
|
// Default to latest run if not specified
|
||||||
state.selectedRunId = task.runs[0].id;
|
state.selectedRunId = task.runs[task.runs.length - 1].id;
|
||||||
} else {
|
} else {
|
||||||
state.selectedRunId = null;
|
state.selectedRunId = null;
|
||||||
}
|
}
|
||||||
|
|
@ -466,20 +495,6 @@ newTaskForm.addEventListener('submit', async (e) => {
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
async function checkSession() {
|
|
||||||
try {
|
|
||||||
const response = await fetch(`${API_URL}/auth/session`, { credentials: 'include' });
|
|
||||||
if (response.ok) {
|
|
||||||
state.isAuthenticated = true;
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
} catch (error) {
|
|
||||||
console.error('Session check failed:', error);
|
|
||||||
}
|
|
||||||
state.isAuthenticated = false;
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
let isPolling = false;
|
let isPolling = false;
|
||||||
async function startAutoRefresh() {
|
async function startAutoRefresh() {
|
||||||
setInterval(async () => {
|
setInterval(async () => {
|
||||||
|
|
@ -499,12 +514,10 @@ async function showLogin() {
|
||||||
}
|
}
|
||||||
|
|
||||||
async function logout() {
|
async function logout() {
|
||||||
try {
|
state.token = null;
|
||||||
await fetch(`${API_URL}/auth/logout`, { method: 'POST', credentials: 'include' });
|
state.refreshToken = null;
|
||||||
} catch (error) {
|
localStorage.removeItem('auth_token');
|
||||||
console.error('Logout failed:', error);
|
localStorage.removeItem('refresh_token');
|
||||||
}
|
|
||||||
state.isAuthenticated = false;
|
|
||||||
showLogin();
|
showLogin();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -518,13 +531,16 @@ async function handleCallback() {
|
||||||
loginOverlay.classList.add('hidden');
|
loginOverlay.classList.add('hidden');
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`, {
|
const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`);
|
||||||
credentials: 'include'
|
|
||||||
});
|
|
||||||
const data = await response.json();
|
const data = await response.json();
|
||||||
|
|
||||||
if (response.ok) {
|
if (data.access_token) {
|
||||||
state.isAuthenticated = true;
|
state.token = data.access_token;
|
||||||
|
localStorage.setItem('auth_token', data.access_token);
|
||||||
|
if (data.refresh_token) {
|
||||||
|
state.refreshToken = data.refresh_token;
|
||||||
|
localStorage.setItem('refresh_token', data.refresh_token);
|
||||||
|
}
|
||||||
callbackOverlay.classList.add('hidden');
|
callbackOverlay.classList.add('hidden');
|
||||||
appEl.classList.remove('hidden');
|
appEl.classList.remove('hidden');
|
||||||
initializeApp();
|
initializeApp();
|
||||||
|
|
@ -548,16 +564,16 @@ logoutBtn.addEventListener('click', () => {
|
||||||
});
|
});
|
||||||
|
|
||||||
async function initializeApp() {
|
async function initializeApp() {
|
||||||
const hasSession = await checkSession();
|
if (!state.token) {
|
||||||
|
|
||||||
if (hasSession) {
|
|
||||||
appEl.classList.remove('hidden');
|
|
||||||
loginOverlay.classList.add('hidden');
|
|
||||||
await fetchTasks();
|
|
||||||
startAutoRefresh();
|
|
||||||
} else {
|
|
||||||
showLogin();
|
showLogin();
|
||||||
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
appEl.classList.remove('hidden');
|
||||||
|
loginOverlay.classList.add('hidden');
|
||||||
|
|
||||||
|
await fetchTasks();
|
||||||
|
startAutoRefresh();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check for callback on load
|
// Check for callback on load
|
||||||
|
|
|
||||||
43
src/agent.rs
43
src/agent.rs
|
|
@ -1,5 +1,4 @@
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
|
|
||||||
use crate::api::{ChatRequest, ChatResponse, Message, Tool};
|
use crate::api::{ChatRequest, ChatResponse, Message, Tool};
|
||||||
use crate::tools;
|
use crate::tools;
|
||||||
|
|
@ -28,6 +27,7 @@ impl Agent {
|
||||||
answer tool once you to give your final answer. current date is {}",
|
answer tool once you to give your final answer. current date is {}",
|
||||||
Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string()
|
Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string()
|
||||||
);
|
);
|
||||||
|
println!("initial_message: {}", intro);
|
||||||
let messages = vec![
|
let messages = vec![
|
||||||
Message {
|
Message {
|
||||||
role: "system".to_string(),
|
role: "system".to_string(),
|
||||||
|
|
@ -62,35 +62,15 @@ impl Agent {
|
||||||
}
|
}
|
||||||
|
|
||||||
fn log(&mut self, message: &str) {
|
fn log(&mut self, message: &str) {
|
||||||
|
println!("{}", message);
|
||||||
self.logs.push_str(message);
|
self.logs.push_str(message);
|
||||||
self.logs.push('\n');
|
self.logs.push('\n');
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn run(&mut self) -> Result<(String, Option<String>), Box<dyn std::error::Error>> {
|
pub async fn run(&mut self) -> Result<(String, Option<String>), Box<dyn std::error::Error>> {
|
||||||
let mut finished = false;
|
let mut file_written = false;
|
||||||
let start_time = Instant::now();
|
|
||||||
let max_duration_secs = std::env::var("AGENT_MAX_DURATION_SECS")
|
|
||||||
.ok()
|
|
||||||
.and_then(|s| s.parse().ok())
|
|
||||||
.unwrap_or(120);
|
|
||||||
let max_duration = Duration::from_secs(max_duration_secs);
|
|
||||||
|
|
||||||
let max_turns = std::env::var("AGENT_MAX_TURNS")
|
while !file_written {
|
||||||
.ok()
|
|
||||||
.and_then(|s| s.parse().ok())
|
|
||||||
.unwrap_or(20);
|
|
||||||
let mut turns = 0;
|
|
||||||
|
|
||||||
while !finished {
|
|
||||||
if start_time.elapsed() > max_duration {
|
|
||||||
return Err("Agent run timed out".into());
|
|
||||||
}
|
|
||||||
|
|
||||||
if turns >= max_turns {
|
|
||||||
return Err("Agent run exceeded max turns".into());
|
|
||||||
}
|
|
||||||
|
|
||||||
turns += 1;
|
|
||||||
let request = ChatRequest {
|
let request = ChatRequest {
|
||||||
model: "kimi-k2.5".to_string(),
|
model: "kimi-k2.5".to_string(),
|
||||||
messages: self.messages.clone(),
|
messages: self.messages.clone(),
|
||||||
|
|
@ -120,12 +100,7 @@ impl Agent {
|
||||||
}
|
}
|
||||||
|
|
||||||
let chat_response: ChatResponse = response.json().await?;
|
let chat_response: ChatResponse = response.json().await?;
|
||||||
let assistant_message = chat_response
|
let assistant_message = chat_response.choices.get(0).unwrap().message.clone();
|
||||||
.choices
|
|
||||||
.get(0)
|
|
||||||
.ok_or("Missing assistant response")?
|
|
||||||
.message
|
|
||||||
.clone();
|
|
||||||
|
|
||||||
self.messages.push(assistant_message.clone());
|
self.messages.push(assistant_message.clone());
|
||||||
|
|
||||||
|
|
@ -153,16 +128,16 @@ impl Agent {
|
||||||
|
|
||||||
self.messages.push(tool_message);
|
self.messages.push(tool_message);
|
||||||
if written {
|
if written {
|
||||||
finished = true;
|
file_written = true;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Continue the loop to send tool results back
|
// Continue the loop to send tool results back
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
// No tool calls from assistant, but we only exit if the task was finished
|
// No more tool calls from assistant, but we only exit if file was written
|
||||||
if !finished {
|
if !file_written {
|
||||||
self.log("--- Assistant didn't finish yet. Waiting for next turn... ---");
|
self.log("--- Assistant didn't use write_file yet. Waiting for next turn... ---");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
43
src/auth.rs
43
src/auth.rs
|
|
@ -10,17 +10,10 @@ pub struct Claims {
|
||||||
pub exp: usize,
|
pub exp: usize,
|
||||||
pub iat: usize,
|
pub iat: usize,
|
||||||
pub iss: String,
|
pub iss: String,
|
||||||
pub aud: Audience,
|
pub aud: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
#[serde(untagged)]
|
|
||||||
pub enum Audience {
|
|
||||||
Single(String),
|
|
||||||
Multiple(Vec<String>),
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Deserialize, Clone)]
|
|
||||||
struct Jwk {
|
struct Jwk {
|
||||||
#[serde(rename = "kty")]
|
#[serde(rename = "kty")]
|
||||||
_kty: String,
|
_kty: String,
|
||||||
|
|
@ -38,17 +31,13 @@ struct Jwks {
|
||||||
|
|
||||||
pub struct JwksVerifier {
|
pub struct JwksVerifier {
|
||||||
issuer: String,
|
issuer: String,
|
||||||
audience: String,
|
|
||||||
jwks_uri: String,
|
jwks_uri: String,
|
||||||
keys: Arc<RwLock<Vec<Jwk>>>,
|
keys: Arc<RwLock<Vec<Jwk>>>,
|
||||||
client: Client,
|
client: Client,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl JwksVerifier {
|
impl JwksVerifier {
|
||||||
pub async fn new(
|
pub async fn new(issuer: String) -> Result<Self, Box<dyn std::error::Error>> {
|
||||||
issuer: String,
|
|
||||||
audience: String,
|
|
||||||
) -> Result<Self, Box<dyn std::error::Error>> {
|
|
||||||
let client = Client::new();
|
let client = Client::new();
|
||||||
// Authentik OIDC discovery
|
// Authentik OIDC discovery
|
||||||
let discovery_url = format!(
|
let discovery_url = format!(
|
||||||
|
|
@ -64,7 +53,6 @@ impl JwksVerifier {
|
||||||
|
|
||||||
let verifier = Self {
|
let verifier = Self {
|
||||||
issuer,
|
issuer,
|
||||||
audience,
|
|
||||||
jwks_uri,
|
jwks_uri,
|
||||||
keys: Arc::new(RwLock::new(Vec::new())),
|
keys: Arc::new(RwLock::new(Vec::new())),
|
||||||
client,
|
client,
|
||||||
|
|
@ -85,29 +73,18 @@ impl JwksVerifier {
|
||||||
let header = decode_header(token)?;
|
let header = decode_header(token)?;
|
||||||
let kid = header.kid.ok_or("Missing kid in token header")?;
|
let kid = header.kid.ok_or("Missing kid in token header")?;
|
||||||
|
|
||||||
let jwk = {
|
let keys = self.keys.read().await;
|
||||||
let keys = self.keys.read().await;
|
let jwk = keys
|
||||||
keys.iter().find(|k| k.kid == kid).cloned()
|
.iter()
|
||||||
};
|
.find(|k| k.kid == kid)
|
||||||
|
.ok_or("Key not found in JWKS")?;
|
||||||
let jwk = match jwk {
|
|
||||||
Some(key) => key,
|
|
||||||
None => {
|
|
||||||
self.refresh_keys().await?;
|
|
||||||
let keys = self.keys.read().await;
|
|
||||||
keys.iter()
|
|
||||||
.find(|k| k.kid == kid)
|
|
||||||
.cloned()
|
|
||||||
.ok_or("Key not found in JWKS")?
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;
|
let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;
|
||||||
|
|
||||||
let mut validation = Validation::new(Algorithm::RS256);
|
let mut validation = Validation::new(Algorithm::RS256);
|
||||||
validation.set_issuer(&[self.issuer.clone()]);
|
validation.set_issuer(&[self.issuer.clone()]);
|
||||||
validation.set_audience(&[self.audience.clone()]);
|
// Aud validation might need careful config, usually it's the client_id
|
||||||
validation.validate_aud = true;
|
validation.validate_aud = false;
|
||||||
|
|
||||||
let token_data = decode::<Claims>(token, &decoding_key, &validation)?;
|
let token_data = decode::<Claims>(token, &decoding_key, &validation)?;
|
||||||
Ok(token_data.claims)
|
Ok(token_data.claims)
|
||||||
|
|
|
||||||
|
|
@ -98,28 +98,19 @@ impl Scheduler {
|
||||||
let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?;
|
let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?;
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let (logs, answer, status) = match agent.run().await {
|
let (logs, answer) = match agent.run().await {
|
||||||
Ok((logs, answer)) => (logs, answer, "completed".to_string()),
|
Ok(res) => res,
|
||||||
Err(e) => (
|
Err(e) => (format!("Scheduled run failed: {}", e), None),
|
||||||
format!("Scheduled run failed: {}", e),
|
|
||||||
None,
|
|
||||||
"failed".to_string(),
|
|
||||||
),
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let run_complete = task_run::ActiveModel {
|
let run_complete = task_run::ActiveModel {
|
||||||
id: Set(run_id),
|
id: Set(run_id),
|
||||||
status: Set(status),
|
status: Set("completed".to_string()),
|
||||||
logs: Set(logs),
|
logs: Set(logs),
|
||||||
answer: Set(answer),
|
answer: Set(answer),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
};
|
};
|
||||||
if let Err(e) = run_complete.update(&db).await {
|
let _ = run_complete.update(&db).await;
|
||||||
eprintln!(
|
|
||||||
"Failed to update scheduled run status for task {}: {}",
|
|
||||||
task_id, e
|
|
||||||
);
|
|
||||||
}
|
|
||||||
});
|
});
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|
|
||||||
221
src/server.rs
221
src/server.rs
|
|
@ -1,13 +1,11 @@
|
||||||
use axum::{
|
use axum::{
|
||||||
Json, RequestPartsExt, Router,
|
Json, RequestPartsExt, Router,
|
||||||
extract::{FromRef, FromRequestParts, Path, Query, State},
|
extract::{FromRef, FromRequestParts, Path, Query, State},
|
||||||
http::{HeaderValue, StatusCode, request::Parts},
|
http::{StatusCode, request::Parts},
|
||||||
response::IntoResponse,
|
|
||||||
routing::{get, post},
|
routing::{get, post},
|
||||||
};
|
};
|
||||||
use axum_extra::{
|
use axum_extra::{
|
||||||
TypedHeader,
|
TypedHeader,
|
||||||
extract::cookie::{Cookie, CookieJar, SameSite},
|
|
||||||
headers::{Authorization, authorization::Bearer},
|
headers::{Authorization, authorization::Bearer},
|
||||||
};
|
};
|
||||||
use chrono::Utc;
|
use chrono::Utc;
|
||||||
|
|
@ -16,7 +14,7 @@ use sea_orm::{
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tower_http::cors::{AllowOrigin, CorsLayer};
|
use tower_http::cors::{Any, CorsLayer};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::agent::Agent;
|
use crate::agent::Agent;
|
||||||
|
|
@ -99,10 +97,7 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
let authentik_client_secret =
|
let authentik_client_secret =
|
||||||
std::env::var("AUTHENTIK_CLIENT_SECRET").map_err(|_| "AUTHENTIK_CLIENT_SECRET not set")?;
|
std::env::var("AUTHENTIK_CLIENT_SECRET").map_err(|_| "AUTHENTIK_CLIENT_SECRET not set")?;
|
||||||
|
|
||||||
let verifier = Arc::new(
|
let verifier = Arc::new(crate::auth::JwksVerifier::new(authentik_issuer.clone()).await?);
|
||||||
crate::auth::JwksVerifier::new(authentik_issuer.clone(), authentik_client_id.clone())
|
|
||||||
.await?,
|
|
||||||
);
|
|
||||||
let authenticator = Arc::new(
|
let authenticator = Arc::new(
|
||||||
crate::auth::Authenticator::new(
|
crate::auth::Authenticator::new(
|
||||||
authentik_issuer,
|
authentik_issuer,
|
||||||
|
|
@ -121,31 +116,19 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
authenticator,
|
authenticator,
|
||||||
});
|
});
|
||||||
|
|
||||||
let cors = build_cors_layer();
|
let cors = CorsLayer::new()
|
||||||
|
.allow_origin(Any)
|
||||||
|
.allow_methods(Any)
|
||||||
|
.allow_headers(Any);
|
||||||
|
|
||||||
let app = Router::new()
|
let app = Router::new()
|
||||||
.route("/api/tasks", post(create_task).get(list_tasks))
|
.route("/api/tasks", post(create_task).get(list_tasks))
|
||||||
.route("/api/tasks/:id", get(get_task).put(update_task))
|
.route("/api/tasks/:id", get(get_task).put(update_task))
|
||||||
.route("/api/tasks/:id/runs", post(rerun_task))
|
.route("/api/tasks/:id/runs", post(rerun_task))
|
||||||
.route("/api/runs/recent", get(get_recent_runs))
|
.route("/api/runs/recent", get(get_recent_runs))
|
||||||
.route("/api/auth/session", get(auth_session))
|
|
||||||
.route("/api/auth/callback", get(auth_callback))
|
.route("/api/auth/callback", get(auth_callback))
|
||||||
.route("/api/auth/refresh", post(auth_refresh))
|
.route("/api/auth/refresh", post(auth_refresh))
|
||||||
.route("/api/auth/logout", post(auth_logout))
|
|
||||||
.layer(cors)
|
.layer(cors)
|
||||||
.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
|
|
||||||
axum::http::header::CONTENT_SECURITY_POLICY,
|
|
||||||
HeaderValue::from_static("default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data:;"),
|
|
||||||
))
|
|
||||||
.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
|
|
||||||
axum::http::header::X_CONTENT_TYPE_OPTIONS,
|
|
||||||
HeaderValue::from_static("nosniff"),
|
|
||||||
))
|
|
||||||
.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
|
|
||||||
axum::http::header::REFERRER_POLICY,
|
|
||||||
HeaderValue::from_static("strict-origin-when-cross-origin"),
|
|
||||||
))
|
|
||||||
.layer(tower_http::limit::RequestBodyLimitLayer::new(1024 * 1024)) // 1MB limit
|
|
||||||
.with_state(state);
|
.with_state(state);
|
||||||
|
|
||||||
let port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string());
|
let port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string());
|
||||||
|
|
@ -157,44 +140,6 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_cors_layer() -> CorsLayer {
|
|
||||||
let origins = std::env::var("CORS_ALLOWED_ORIGINS").ok();
|
|
||||||
|
|
||||||
let allow_origin = if let Some(origins) = origins {
|
|
||||||
let values: Vec<HeaderValue> = origins
|
|
||||||
.split(',')
|
|
||||||
.map(|origin| origin.trim())
|
|
||||||
.filter(|origin| !origin.is_empty())
|
|
||||||
.filter_map(|origin| HeaderValue::from_str(origin).ok())
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
if values.is_empty() {
|
|
||||||
AllowOrigin::mirror_request()
|
|
||||||
} else {
|
|
||||||
AllowOrigin::list(values)
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
AllowOrigin::mirror_request()
|
|
||||||
};
|
|
||||||
|
|
||||||
CorsLayer::new()
|
|
||||||
.allow_origin(allow_origin)
|
|
||||||
.allow_methods([
|
|
||||||
axum::http::Method::GET,
|
|
||||||
axum::http::Method::POST,
|
|
||||||
axum::http::Method::PUT,
|
|
||||||
axum::http::Method::PATCH,
|
|
||||||
axum::http::Method::DELETE,
|
|
||||||
axum::http::Method::OPTIONS,
|
|
||||||
])
|
|
||||||
.allow_headers([
|
|
||||||
axum::http::header::CONTENT_TYPE,
|
|
||||||
axum::http::header::AUTHORIZATION,
|
|
||||||
axum::http::header::ACCEPT,
|
|
||||||
])
|
|
||||||
.allow_credentials(true)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn list_tasks(
|
async fn list_tasks(
|
||||||
_user: AuthenticatedUser,
|
_user: AuthenticatedUser,
|
||||||
State(state): State<Arc<AppState>>,
|
State(state): State<Arc<AppState>>,
|
||||||
|
|
@ -304,13 +249,9 @@ async fn execute_agent_run(
|
||||||
)
|
)
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
|
|
||||||
let (logs, answer, status) = match agent.run().await {
|
let (logs, answer) = match agent.run().await {
|
||||||
Ok((logs, answer)) => (logs, answer, "completed".to_string()),
|
Ok((logs, answer)) => (logs, answer),
|
||||||
Err(e) => (
|
Err(e) => (format!("Execution failed: {}", e), None),
|
||||||
format!("Execution failed: {}", e),
|
|
||||||
None,
|
|
||||||
"failed".to_string(),
|
|
||||||
),
|
|
||||||
};
|
};
|
||||||
|
|
||||||
// Update with final logs and status
|
// Update with final logs and status
|
||||||
|
|
@ -326,7 +267,7 @@ async fn execute_agent_run(
|
||||||
|
|
||||||
run.logs = Set(logs.clone());
|
run.logs = Set(logs.clone());
|
||||||
run.answer = Set(answer.clone());
|
run.answer = Set(answer.clone());
|
||||||
run.status = Set(status);
|
run.status = Set("completed".to_string());
|
||||||
|
|
||||||
run.update(&state.db)
|
run.update(&state.db)
|
||||||
.await
|
.await
|
||||||
|
|
@ -378,27 +319,26 @@ where
|
||||||
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
|
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
|
||||||
let app_state = Arc::<AppState>::from_ref(state);
|
let app_state = Arc::<AppState>::from_ref(state);
|
||||||
|
|
||||||
let token = if let Ok(TypedHeader(Authorization(bearer))) =
|
let TypedHeader(Authorization(bearer)) = parts
|
||||||
parts.extract::<TypedHeader<Authorization<Bearer>>>().await
|
.extract::<TypedHeader<Authorization<Bearer>>>()
|
||||||
{
|
.await
|
||||||
Some(bearer.token().to_string())
|
.map_err(|_| {
|
||||||
} else {
|
(
|
||||||
let jar = parts.extract::<CookieJar>().await.unwrap();
|
StatusCode::UNAUTHORIZED,
|
||||||
jar.get("access_token")
|
"Missing or invalid Authorization header".to_string(),
|
||||||
.map(|cookie| cookie.value().to_string())
|
)
|
||||||
};
|
})?;
|
||||||
|
|
||||||
let token = token.ok_or((
|
let claims = app_state
|
||||||
StatusCode::UNAUTHORIZED,
|
.verifier
|
||||||
"Missing or invalid access token".to_string(),
|
.verify(bearer.token())
|
||||||
))?;
|
.await
|
||||||
|
.map_err(|e| {
|
||||||
let claims = app_state.verifier.verify(&token).await.map_err(|e| {
|
(
|
||||||
(
|
StatusCode::UNAUTHORIZED,
|
||||||
StatusCode::UNAUTHORIZED,
|
format!("Token verification failed: {}", e),
|
||||||
format!("Token verification failed: {}", e),
|
)
|
||||||
)
|
})?;
|
||||||
})?;
|
|
||||||
|
|
||||||
Ok(AuthenticatedUser(claims))
|
Ok(AuthenticatedUser(claims))
|
||||||
}
|
}
|
||||||
|
|
@ -412,117 +352,36 @@ pub struct AuthCallbackQuery {
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
struct RefreshRequest {
|
struct RefreshRequest {
|
||||||
refresh_token: Option<String>,
|
refresh_token: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn auth_refresh(
|
async fn auth_refresh(
|
||||||
State(state): State<Arc<AppState>>,
|
State(state): State<Arc<AppState>>,
|
||||||
jar: CookieJar,
|
|
||||||
Json(payload): Json<RefreshRequest>,
|
Json(payload): Json<RefreshRequest>,
|
||||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let refresh_token = payload
|
state
|
||||||
.refresh_token
|
|
||||||
.filter(|token| !token.is_empty())
|
|
||||||
.or_else(|| {
|
|
||||||
jar.get("refresh_token")
|
|
||||||
.map(|cookie| cookie.value().to_string())
|
|
||||||
})
|
|
||||||
.ok_or((
|
|
||||||
StatusCode::UNAUTHORIZED,
|
|
||||||
"Missing refresh token".to_string(),
|
|
||||||
))?;
|
|
||||||
|
|
||||||
let data = state
|
|
||||||
.authenticator
|
.authenticator
|
||||||
.refresh_token(refresh_token)
|
.refresh_token(payload.refresh_token)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?;
|
.map(Json)
|
||||||
|
.map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))
|
||||||
let jar = update_auth_cookies(jar, &data);
|
|
||||||
Ok((jar, Json(data)))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn auth_callback(
|
async fn auth_callback(
|
||||||
State(state): State<Arc<AppState>>,
|
State(state): State<Arc<AppState>>,
|
||||||
jar: CookieJar,
|
|
||||||
Query(query): Query<AuthCallbackQuery>,
|
Query(query): Query<AuthCallbackQuery>,
|
||||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||||
let data = state
|
state
|
||||||
.authenticator
|
.authenticator
|
||||||
.exchange_code(query.code, query.redirect_uri)
|
.exchange_code(query.code, query.redirect_uri)
|
||||||
.await
|
.await
|
||||||
|
.map(Json)
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
(
|
(
|
||||||
StatusCode::INTERNAL_SERVER_ERROR,
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
format!("Token exchange failed: {}", e),
|
format!("Token exchange failed: {}", e),
|
||||||
)
|
)
|
||||||
})?;
|
})
|
||||||
|
|
||||||
let jar = update_auth_cookies(jar, &data);
|
|
||||||
Ok((jar, Json(data)))
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn auth_logout(jar: CookieJar) -> impl IntoResponse {
|
|
||||||
let jar = clear_auth_cookies(jar);
|
|
||||||
(jar, StatusCode::NO_CONTENT)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn auth_session(user: AuthenticatedUser) -> Json<serde_json::Value> {
|
|
||||||
Json(serde_json::json!({
|
|
||||||
"authenticated": true,
|
|
||||||
"user": user.0
|
|
||||||
}))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn secure() -> bool {
|
|
||||||
std::env::var("COOKIE_SECURE")
|
|
||||||
.map(|value| value == "true")
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn update_auth_cookies(jar: CookieJar, data: &serde_json::Value) -> CookieJar {
|
|
||||||
let access_token = data.get("access_token");
|
|
||||||
let refresh_token = data.get("refresh_token");
|
|
||||||
|
|
||||||
let mut jar = jar;
|
|
||||||
|
|
||||||
if let Some(token) = access_token.and_then(|t| t.as_str()) {
|
|
||||||
let cookie = Cookie::build(("access_token", token.to_owned()))
|
|
||||||
.path("/")
|
|
||||||
.http_only(true)
|
|
||||||
.same_site(SameSite::Lax)
|
|
||||||
.secure(secure())
|
|
||||||
.build();
|
|
||||||
jar = jar.add(cookie);
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(token) = refresh_token.and_then(|t| t.as_str()) {
|
|
||||||
let cookie = Cookie::build(("refresh_token", token.to_owned()))
|
|
||||||
.path("/")
|
|
||||||
.http_only(true)
|
|
||||||
.same_site(SameSite::Lax)
|
|
||||||
.secure(secure())
|
|
||||||
.build();
|
|
||||||
jar = jar.add(cookie);
|
|
||||||
}
|
|
||||||
|
|
||||||
jar
|
|
||||||
}
|
|
||||||
|
|
||||||
fn clear_auth_cookies(jar: CookieJar) -> CookieJar {
|
|
||||||
let mut jar = jar;
|
|
||||||
for name in ["access_token", "refresh_token"] {
|
|
||||||
let cookie = Cookie::build((name, ""))
|
|
||||||
.path("/")
|
|
||||||
.http_only(true)
|
|
||||||
.same_site(SameSite::Lax)
|
|
||||||
.secure(secure())
|
|
||||||
.max_age(cookie::time::Duration::seconds(0))
|
|
||||||
.build();
|
|
||||||
jar = jar.add(cookie);
|
|
||||||
}
|
|
||||||
|
|
||||||
jar
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn get_task(
|
async fn get_task(
|
||||||
|
|
|
||||||
|
|
@ -51,6 +51,11 @@ pub async fn handle_tool_call(
|
||||||
let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
|
let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
|
||||||
let query = args.get("query").ok_or("Missing query argument")?;
|
let query = args.get("query").ok_or("Missing query argument")?;
|
||||||
|
|
||||||
|
println!(
|
||||||
|
"--- Executing tool: google_search(query: \"{}\") ---",
|
||||||
|
query
|
||||||
|
);
|
||||||
|
|
||||||
let search_result = if let Some(key) = tavily_api_key {
|
let search_result = if let Some(key) = tavily_api_key {
|
||||||
match api::perform_search(query, key).await {
|
match api::perform_search(query, key).await {
|
||||||
Ok(results) => results,
|
Ok(results) => results,
|
||||||
|
|
@ -64,6 +69,8 @@ pub async fn handle_tool_call(
|
||||||
let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
|
let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
|
||||||
let result = args.get("result").ok_or("Missing result argument")?;
|
let result = args.get("result").ok_or("Missing result argument")?;
|
||||||
|
|
||||||
|
println!("--- Finishing task: {}", result);
|
||||||
|
|
||||||
answer = Some(result.clone());
|
answer = Some(result.clone());
|
||||||
(result.clone(), true)
|
(result.clone(), true)
|
||||||
} else {
|
} else {
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue