This commit is contained in:
parent
cc8fcd682b
commit
91db5861c8
21 changed files with 1263 additions and 79 deletions
|
|
@ -14,6 +14,8 @@ pub struct Config {
|
|||
pub cookie_secure: bool,
|
||||
pub agent_max_turns: u32,
|
||||
pub agent_max_duration_secs: u64,
|
||||
pub vapid_private_key: String,
|
||||
pub vapid_public_key: String,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
|
|
@ -54,6 +56,11 @@ impl Config {
|
|||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(120);
|
||||
|
||||
let vapid_private_key = env::var("VAPID_PRIVATE_KEY")
|
||||
.map_err(|_| AppError::Config("VAPID_PRIVATE_KEY must be set".into()))?;
|
||||
let vapid_public_key = env::var("VAPID_PUBLIC_KEY")
|
||||
.map_err(|_| AppError::Config("VAPID_PUBLIC_KEY must be set".into()))?;
|
||||
|
||||
Ok(Config {
|
||||
database_url,
|
||||
port,
|
||||
|
|
@ -66,6 +73,8 @@ impl Config {
|
|||
cookie_secure,
|
||||
agent_max_turns,
|
||||
agent_max_duration_secs,
|
||||
vapid_private_key,
|
||||
vapid_public_key,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
pub mod agent;
|
||||
pub mod auth;
|
||||
pub mod notifications;
|
||||
pub mod tasks;
|
||||
|
|
|
|||
1
src/domain/notifications/mod.rs
Normal file
1
src/domain/notifications/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod push;
|
||||
101
src/domain/notifications/push.rs
Normal file
101
src/domain/notifications/push.rs
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
use crate::error::AppResult;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use web_push::*;
|
||||
|
||||
pub struct PushSender {
|
||||
pub private_key: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone, Debug)]
|
||||
pub struct PushSubscription {
|
||||
pub endpoint: String,
|
||||
pub p256dh: String,
|
||||
pub auth: String,
|
||||
}
|
||||
|
||||
fn ensure_pem(input: &str) -> String {
|
||||
let input = input.trim();
|
||||
if input.contains("-----BEGIN") {
|
||||
return input.to_string();
|
||||
}
|
||||
|
||||
if input.starts_with("MHc") {
|
||||
format!(
|
||||
"-----BEGIN EC PRIVATE KEY-----\n{}\n-----END EC PRIVATE KEY-----",
|
||||
input
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"-----BEGIN PRIVATE KEY-----\n{}\n-----END PRIVATE KEY-----",
|
||||
input
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl PushSender {
|
||||
pub fn new(private_key_pem: &str) -> AppResult<Self> {
|
||||
let pem = ensure_pem(private_key_pem);
|
||||
// Validate key immediately to catch config errors early
|
||||
let _ =
|
||||
VapidSignatureBuilder::from_pem_no_sub(std::io::Cursor::new(&pem)).map_err(|e| {
|
||||
crate::error::AppError::Internal(format!("Invalid VAPID private key: {}", e))
|
||||
})?;
|
||||
|
||||
Ok(Self { private_key: pem })
|
||||
}
|
||||
|
||||
pub fn get_public_key(&self) -> AppResult<Vec<u8>> {
|
||||
let builder =
|
||||
VapidSignatureBuilder::from_pem_no_sub(std::io::Cursor::new(&self.private_key))
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
Ok(builder.get_public_key())
|
||||
}
|
||||
|
||||
pub async fn send_notification(
|
||||
&self,
|
||||
subscription: &PushSubscription,
|
||||
title: &str,
|
||||
body: &str,
|
||||
) -> AppResult<()> {
|
||||
let subscription_info = SubscriptionInfo::new(
|
||||
subscription.endpoint.clone(),
|
||||
subscription.p256dh.clone(),
|
||||
subscription.auth.clone(),
|
||||
);
|
||||
|
||||
let builder =
|
||||
VapidSignatureBuilder::from_pem_no_sub(std::io::Cursor::new(&self.private_key))
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
let vapid_signature = builder
|
||||
.add_sub_info(&subscription_info)
|
||||
.build()
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
let mut builder = WebPushMessageBuilder::new(&subscription_info);
|
||||
|
||||
builder.set_vapid_signature(vapid_signature);
|
||||
|
||||
let payload = serde_json::to_vec(&serde_json::json!({
|
||||
"title": title,
|
||||
"body": body,
|
||||
}))
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
builder.set_payload(ContentEncoding::Aes128Gcm, &payload);
|
||||
|
||||
let message = builder
|
||||
.build()
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
let client = IsahcWebPushClient::new()
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
client
|
||||
.send(message)
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
use chrono::Utc;
|
||||
use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait, Set};
|
||||
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
|
@ -8,6 +8,7 @@ use crate::config::Config;
|
|||
use crate::domain::agent::Agent;
|
||||
use crate::entities::task::Entity as Task;
|
||||
use crate::entities::task_run::{self, Entity as TaskRun};
|
||||
use crate::entities::{push_subscription, task_subscription};
|
||||
use crate::scheduler::Scheduler;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
|
|
@ -119,7 +120,7 @@ pub async fn execute_agent_run(
|
|||
let mut run = run;
|
||||
run.logs = Set(logs.clone());
|
||||
run.answer = Set(answer.clone());
|
||||
run.status = Set(status);
|
||||
run.status = Set(status.clone());
|
||||
|
||||
run.update(db)
|
||||
.await
|
||||
|
|
@ -132,6 +133,38 @@ pub async fn execute_agent_run(
|
|||
task_response.clone(),
|
||||
));
|
||||
|
||||
// Send Push Notifications to subscribers
|
||||
let subscriptions = task_subscription::Entity::find()
|
||||
.filter(task_subscription::Column::TaskId.eq(task_id))
|
||||
.all(db)
|
||||
.await
|
||||
.map_err(crate::error::AppError::Database)?;
|
||||
|
||||
for sub in subscriptions {
|
||||
let push_subs = push_subscription::Entity::find()
|
||||
.filter(push_subscription::Column::UserSub.eq(sub.user_sub))
|
||||
.all(db)
|
||||
.await
|
||||
.map_err(crate::error::AppError::Database)?;
|
||||
|
||||
for push_sub in push_subs {
|
||||
let sender = _scheduler.push_sender.clone();
|
||||
let goal = task_response.goal.clone();
|
||||
let status = status.clone();
|
||||
let sub_data = crate::domain::notifications::push::PushSubscription {
|
||||
endpoint: push_sub.endpoint,
|
||||
p256dh: push_sub.p256dh,
|
||||
auth: push_sub.auth,
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _ = sender
|
||||
.send_notification(&sub_data, &format!("Task Completed: {}", status), &goal)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(task_response)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
pub mod push_subscription;
|
||||
pub mod task;
|
||||
pub mod task_run;
|
||||
pub mod task_subscription;
|
||||
|
|
|
|||
19
src/entities/push_subscription.rs
Normal file
19
src/entities/push_subscription.rs
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
use sea_orm::entity::prelude::*;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
|
||||
#[sea_orm(table_name = "push_subscriptions")]
|
||||
pub struct Model {
|
||||
#[sea_orm(primary_key, auto_increment = false)]
|
||||
pub id: Uuid,
|
||||
pub user_sub: String,
|
||||
pub endpoint: String,
|
||||
pub p256dh: String,
|
||||
pub auth: String,
|
||||
pub created_at: DateTimeWithTimeZone,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
|
||||
pub enum Relation {}
|
||||
|
||||
impl ActiveModelBehavior for ActiveModel {}
|
||||
32
src/entities/task_subscription.rs
Normal file
32
src/entities/task_subscription.rs
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
use sea_orm::entity::prelude::*;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
|
||||
#[sea_orm(table_name = "task_subscriptions")]
|
||||
pub struct Model {
|
||||
#[sea_orm(primary_key, auto_increment = false)]
|
||||
pub id: Uuid,
|
||||
pub user_sub: String,
|
||||
pub task_id: Uuid,
|
||||
pub created_at: DateTimeWithTimeZone,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
|
||||
pub enum Relation {
|
||||
#[sea_orm(
|
||||
belongs_to = "super::task::Entity",
|
||||
from = "Column::TaskId",
|
||||
to = "super::task::Column::Id",
|
||||
on_update = "NoAction",
|
||||
on_delete = "Cascade"
|
||||
)]
|
||||
Task,
|
||||
}
|
||||
|
||||
impl Related<super::task::Entity> for Entity {
|
||||
fn to() -> RelationDef {
|
||||
Relation::Task.def()
|
||||
}
|
||||
}
|
||||
|
||||
impl ActiveModelBehavior for ActiveModel {}
|
||||
|
|
@ -14,6 +14,7 @@ pub struct Scheduler {
|
|||
tasks_to_jobs: DashMap<Uuid, Uuid>,
|
||||
config: Arc<crate::config::Config>,
|
||||
pub tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
|
||||
pub push_sender: Arc<crate::domain::notifications::push::PushSender>,
|
||||
}
|
||||
|
||||
impl Scheduler {
|
||||
|
|
@ -21,6 +22,7 @@ impl Scheduler {
|
|||
db: DatabaseConnection,
|
||||
config: Arc<crate::config::Config>,
|
||||
tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
|
||||
push_sender: Arc<crate::domain::notifications::push::PushSender>,
|
||||
) -> AppResult<Self> {
|
||||
let scheduler = JobScheduler::new()
|
||||
.await
|
||||
|
|
@ -35,6 +37,7 @@ impl Scheduler {
|
|||
tasks_to_jobs: DashMap::new(),
|
||||
config,
|
||||
tx,
|
||||
push_sender,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -47,13 +50,15 @@ impl Scheduler {
|
|||
let db = self.db.clone();
|
||||
let config = self.config.clone();
|
||||
let tx = self.tx.clone();
|
||||
let push_sender = self.push_sender.clone();
|
||||
|
||||
let job = Job::new_async(cron_expr, move |_uuid, _l| {
|
||||
let db = db.clone();
|
||||
let config = config.clone();
|
||||
let tx = tx.clone();
|
||||
let push_sender = push_sender.clone();
|
||||
Box::pin(async move {
|
||||
if let Err(e) = Self::run_task(db, config, tx, task_id).await {
|
||||
if let Err(e) = Self::run_task(db, config, tx, push_sender, task_id).await {
|
||||
tracing::error!("Error in scheduled task {}: {}", task_id, e);
|
||||
}
|
||||
})
|
||||
|
|
@ -87,6 +92,7 @@ impl Scheduler {
|
|||
db: DatabaseConnection,
|
||||
config: Arc<crate::config::Config>,
|
||||
tx: tokio::sync::broadcast::Sender<crate::server::notifications::WsEvent>,
|
||||
push_sender: Arc<crate::domain::notifications::push::PushSender>,
|
||||
task_id: Uuid,
|
||||
) -> AppResult<()> {
|
||||
let task = Task::find_by_id(task_id)
|
||||
|
|
@ -147,7 +153,7 @@ impl Scheduler {
|
|||
|
||||
let run_complete = task_run::ActiveModel {
|
||||
id: Set(run_id),
|
||||
status: Set(status),
|
||||
status: Set(status.clone()),
|
||||
logs: Set(logs),
|
||||
answer: Set(answer),
|
||||
..Default::default()
|
||||
|
|
@ -162,8 +168,46 @@ impl Scheduler {
|
|||
|
||||
if let Ok(task_response) = crate::domain::tasks::get_task_inner(task_id, &db).await {
|
||||
let _ = tx.send(crate::server::notifications::WsEvent::RunFinished(
|
||||
task_response,
|
||||
task_response.clone(),
|
||||
));
|
||||
|
||||
// Push notifications for scheduled runs
|
||||
use crate::entities::{push_subscription, task_subscription};
|
||||
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
|
||||
if let Ok(subscriptions) = task_subscription::Entity::find()
|
||||
.filter(task_subscription::Column::TaskId.eq(task_id))
|
||||
.all(&db)
|
||||
.await
|
||||
{
|
||||
for sub in subscriptions {
|
||||
if let Ok(push_subs) = push_subscription::Entity::find()
|
||||
.filter(push_subscription::Column::UserSub.eq(sub.user_sub))
|
||||
.all(&db)
|
||||
.await
|
||||
{
|
||||
for push_sub in push_subs {
|
||||
let sender = push_sender.clone();
|
||||
let sub_data =
|
||||
crate::domain::notifications::push::PushSubscription {
|
||||
endpoint: push_sub.endpoint,
|
||||
p256dh: push_sub.p256dh,
|
||||
auth: push_sub.auth,
|
||||
};
|
||||
let goal = task_response.goal.clone();
|
||||
let status = status.clone();
|
||||
tokio::spawn(async move {
|
||||
let _ = sender
|
||||
.send_notification(
|
||||
&sub_data,
|
||||
&format!("Scheduled Task Completed: {}", status),
|
||||
&goal,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -35,8 +35,12 @@ pub async fn start(config: crate::config::Config) -> AppResult<()> {
|
|||
|
||||
let (tx, _) = tokio::sync::broadcast::channel(100);
|
||||
|
||||
let push_sender = Arc::new(crate::domain::notifications::push::PushSender::new(
|
||||
&config.vapid_private_key.clone(),
|
||||
)?);
|
||||
|
||||
let scheduler = Arc::new(
|
||||
Scheduler::new(db.clone(), config.clone(), tx.clone())
|
||||
Scheduler::new(db.clone(), config.clone(), tx.clone(), push_sender.clone())
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?,
|
||||
);
|
||||
|
|
@ -128,6 +132,10 @@ fn build_app(state: Arc<AppState>, config: &crate::config::Config) -> Router {
|
|||
.route("/api/auth/logout", post(auth::auth_logout))
|
||||
.route("/api/chat", post(chat::chat_handler))
|
||||
.route("/api/ws", get(notifications::ws_handler))
|
||||
.route("/api/notifications/register", post(notifications::push_handlers::register_push))
|
||||
.route("/api/notifications/vapid-key", get(notifications::push_handlers::get_vapid_key))
|
||||
.route("/api/tasks/:id/subscription", get(notifications::push_handlers::get_subscription_status))
|
||||
.route("/api/tasks/:id/subscribe", post(notifications::push_handlers::subscribe_task).delete(notifications::push_handlers::unsubscribe_task))
|
||||
.layer(cors)
|
||||
.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
|
||||
axum::http::header::CONTENT_SECURITY_POLICY,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
pub mod push_handlers;
|
||||
|
||||
use crate::domain::tasks::{RecentRunResponse, TaskResponse};
|
||||
use crate::server::AppState;
|
||||
use axum::{
|
||||
|
|
|
|||
128
src/server/notifications/push_handlers.rs
Normal file
128
src/server/notifications/push_handlers.rs
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
use axum::{
|
||||
Json,
|
||||
extract::{Path, State},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, Set};
|
||||
use std::sync::Arc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::entities::{push_subscription, task_subscription};
|
||||
use crate::error::{AppError, AppResult};
|
||||
use crate::server::AppState;
|
||||
use crate::server::auth::AuthenticatedUser;
|
||||
use base64::Engine;
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub struct RegisterPushRequest {
|
||||
pub endpoint: String,
|
||||
pub p256dh: String,
|
||||
pub auth: String,
|
||||
}
|
||||
|
||||
pub async fn register_push(
|
||||
State(state): State<Arc<AppState>>,
|
||||
user: AuthenticatedUser,
|
||||
Json(payload): Json<RegisterPushRequest>,
|
||||
) -> AppResult<Json<serde_json::Value>> {
|
||||
let user_sub = user.0.sub;
|
||||
|
||||
// Check if subscription exists
|
||||
let existing = push_subscription::Entity::find()
|
||||
.filter(push_subscription::Column::UserSub.eq(user_sub.clone()))
|
||||
.filter(push_subscription::Column::Endpoint.eq(payload.endpoint.clone()))
|
||||
.one(&state.db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
|
||||
if existing.is_none() {
|
||||
let new_sub = push_subscription::ActiveModel {
|
||||
id: Set(Uuid::new_v4()),
|
||||
user_sub: Set(user_sub),
|
||||
endpoint: Set(payload.endpoint),
|
||||
p256dh: Set(payload.p256dh),
|
||||
auth: Set(payload.auth),
|
||||
created_at: Set(Utc::now().into()),
|
||||
};
|
||||
new_sub
|
||||
.insert(&state.db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
}
|
||||
|
||||
Ok(Json(serde_json::json!({ "status": "registered" })))
|
||||
}
|
||||
|
||||
pub async fn subscribe_task(
|
||||
State(state): State<Arc<AppState>>,
|
||||
user: AuthenticatedUser,
|
||||
Path(task_id): Path<Uuid>,
|
||||
) -> AppResult<Json<serde_json::Value>> {
|
||||
let user_sub = user.0.sub;
|
||||
|
||||
let existing = task_subscription::Entity::find()
|
||||
.filter(task_subscription::Column::UserSub.eq(user_sub.clone()))
|
||||
.filter(task_subscription::Column::TaskId.eq(task_id))
|
||||
.one(&state.db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
|
||||
if existing.is_none() {
|
||||
let new_sub = task_subscription::ActiveModel {
|
||||
id: Set(Uuid::new_v4()),
|
||||
user_sub: Set(user_sub),
|
||||
task_id: Set(task_id),
|
||||
created_at: Set(Utc::now().into()),
|
||||
};
|
||||
new_sub
|
||||
.insert(&state.db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
}
|
||||
|
||||
Ok(Json(serde_json::json!({ "status": "subscribed" })))
|
||||
}
|
||||
|
||||
pub async fn unsubscribe_task(
|
||||
State(state): State<Arc<AppState>>,
|
||||
user: AuthenticatedUser,
|
||||
Path(task_id): Path<Uuid>,
|
||||
) -> AppResult<Json<serde_json::Value>> {
|
||||
let user_sub = user.0.sub;
|
||||
|
||||
task_subscription::Entity::delete_many()
|
||||
.filter(task_subscription::Column::UserSub.eq(user_sub))
|
||||
.filter(task_subscription::Column::TaskId.eq(task_id))
|
||||
.exec(&state.db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
|
||||
Ok(Json(serde_json::json!({ "status": "unsubscribed" })))
|
||||
}
|
||||
|
||||
pub async fn get_vapid_key(
|
||||
State(state): State<Arc<AppState>>,
|
||||
) -> AppResult<Json<serde_json::Value>> {
|
||||
let public_key = state.scheduler.push_sender.get_public_key()?;
|
||||
let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(public_key);
|
||||
Ok(Json(serde_json::json!({ "publicKey": encoded })))
|
||||
}
|
||||
|
||||
pub async fn get_subscription_status(
|
||||
State(state): State<Arc<AppState>>,
|
||||
user: AuthenticatedUser,
|
||||
Path(task_id): Path<Uuid>,
|
||||
) -> AppResult<Json<serde_json::Value>> {
|
||||
let user_sub = user.0.sub;
|
||||
|
||||
let existing = task_subscription::Entity::find()
|
||||
.filter(task_subscription::Column::UserSub.eq(user_sub))
|
||||
.filter(task_subscription::Column::TaskId.eq(task_id))
|
||||
.one(&state.db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
|
||||
Ok(Json(
|
||||
serde_json::json!({ "isSubscribed": existing.is_some() }),
|
||||
))
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue