This commit is contained in:
parent
cc8fcd682b
commit
91db5861c8
21 changed files with 1263 additions and 79 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -1,3 +1,4 @@
|
|||
/target
|
||||
node_modules
|
||||
/frontend/dist
|
||||
/frontend/dist
|
||||
.env
|
||||
691
Cargo.lock
generated
691
Cargo.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -22,6 +22,8 @@ tokio-cron-scheduler = "0.15.1"
|
|||
dashmap = "6.1.0"
|
||||
jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] }
|
||||
base64 = "0.22.1"
|
||||
web-push = { version = "0.10.0", features = ["isahc-client"] }
|
||||
isahc = "1.7"
|
||||
axum-extra = { version = "0.9", features = ["typed-header", "cookie"] }
|
||||
cookie = "0.18"
|
||||
thiserror = "2.0.18"
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@
|
|||
<link rel="apple-touch-icon" href="/apple-touch-icon.png">
|
||||
<link rel="manifest" href="/manifest.json">
|
||||
<title>Antigravity Agent Dashboard</title>
|
||||
<link rel="stylesheet" href="/src/style.css">
|
||||
<link rel="preconnect" href="https://fonts.googleapis.com">
|
||||
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
|
||||
<link
|
||||
|
|
@ -145,6 +144,9 @@
|
|||
<button id="rerun-btn" class="btn btn-ghost">
|
||||
<span>↻</span> Re-run Task
|
||||
</button>
|
||||
<button id="notify-task-btn" class="btn btn-ghost">
|
||||
<span class="notify-icon">🔔</span> Notify Me
|
||||
</button>
|
||||
</div>
|
||||
</header>
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import { marked } from 'marked';
|
||||
import DOMPurify from 'dompurify';
|
||||
import './style.css';
|
||||
|
||||
const API_URL = '/api';
|
||||
// These should ideally be environment-specific
|
||||
|
|
@ -244,6 +245,19 @@ function selectTask(id, runId = null) {
|
|||
renderRunHistory(task);
|
||||
showTaskView(task);
|
||||
|
||||
// Fetch and update subscription status
|
||||
fetch(`${API_URL}/tasks/${id}/subscription`)
|
||||
.then(res => res.json())
|
||||
.then(data => {
|
||||
const btn = document.getElementById('notify-task-btn');
|
||||
if (data.isSubscribed) {
|
||||
btn.classList.add('notified');
|
||||
} else {
|
||||
btn.classList.remove('notified');
|
||||
}
|
||||
})
|
||||
.catch(err => console.error('Failed to fetch subscription status', err));
|
||||
|
||||
// Close sidebar on mobile after selection
|
||||
if (window.innerWidth <= 768) {
|
||||
closeMobileMenu();
|
||||
|
|
@ -802,7 +816,68 @@ if (window.location.pathname === '/callback' || window.location.search.includes(
|
|||
if ('serviceWorker' in navigator) {
|
||||
window.addEventListener('load', () => {
|
||||
navigator.serviceWorker.register('/sw.js')
|
||||
.then(reg => console.log('SW registered', reg))
|
||||
.then(reg => {
|
||||
console.log('SW registered', reg);
|
||||
setupPush(reg);
|
||||
})
|
||||
.catch(err => console.error('SW registration failed', err));
|
||||
});
|
||||
}
|
||||
|
||||
async function setupPush(registration) {
|
||||
try {
|
||||
const vapidResponse = await fetch(`${API_URL}/notifications/vapid-key`);
|
||||
const { publicKey } = await vapidResponse.json();
|
||||
|
||||
const subscription = await registration.pushManager.subscribe({
|
||||
userVisibleOnly: true,
|
||||
applicationServerKey: urlBase64ToUint8Array(publicKey)
|
||||
});
|
||||
|
||||
await fetch(`${API_URL}/notifications/register`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
endpoint: subscription.endpoint,
|
||||
p256dh: b64(subscription.getKey('p256dh')),
|
||||
auth: b64(subscription.getKey('auth'))
|
||||
})
|
||||
});
|
||||
console.log('Push registered');
|
||||
} catch (err) {
|
||||
console.warn('Push registration failed:', err);
|
||||
}
|
||||
}
|
||||
|
||||
function b64(buffer) {
|
||||
return btoa(String.fromCharCode.apply(null, new Uint8Array(buffer)));
|
||||
}
|
||||
|
||||
function urlBase64ToUint8Array(base64String) {
|
||||
const padding = '='.repeat((4 - base64String.length % 4) % 4);
|
||||
const base64 = (base64String + padding).replace(/\-/g, '+').replace(/_/g, '/');
|
||||
const rawData = window.atob(base64);
|
||||
const outputArray = new Uint8Array(rawData.length);
|
||||
for (let i = 0; i < rawData.length; ++i) {
|
||||
outputArray[i] = rawData.charCodeAt(i);
|
||||
}
|
||||
return outputArray;
|
||||
}
|
||||
|
||||
async function toggleTaskSubscription(taskId) {
|
||||
const btn = document.getElementById('notify-task-btn');
|
||||
const isNotified = btn.classList.contains('notified');
|
||||
const method = isNotified ? 'DELETE' : 'POST';
|
||||
|
||||
try {
|
||||
await fetch(`${API_URL}/tasks/${taskId}/subscribe`, { method });
|
||||
btn.classList.toggle('notified');
|
||||
showToast(isNotified ? 'Notifications disabled' : 'Notifications enabled');
|
||||
} catch (err) {
|
||||
showToast('Failed to update notifications');
|
||||
}
|
||||
}
|
||||
|
||||
document.getElementById('notify-task-btn').addEventListener('click', () => {
|
||||
if (state.selectedTaskId) toggleTaskSubscription(state.selectedTaskId);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1101,6 +1101,37 @@ textarea:focus {
|
|||
}
|
||||
|
||||
/* Responsive Styles */
|
||||
#notify-task-btn.notified {
|
||||
color: var(--primary);
|
||||
background: var(--primary-glow);
|
||||
}
|
||||
|
||||
#notify-task-btn.notified .notify-icon {
|
||||
animation: ring 0.5s ease;
|
||||
}
|
||||
|
||||
@keyframes ring {
|
||||
0% {
|
||||
transform: rotate(0);
|
||||
}
|
||||
|
||||
25% {
|
||||
transform: rotate(15deg);
|
||||
}
|
||||
|
||||
50% {
|
||||
transform: rotate(-15deg);
|
||||
}
|
||||
|
||||
75% {
|
||||
transform: rotate(10deg);
|
||||
}
|
||||
|
||||
100% {
|
||||
transform: rotate(0);
|
||||
}
|
||||
}
|
||||
|
||||
@media (max-width: 768px) {
|
||||
.menu-toggle {
|
||||
display: flex;
|
||||
|
|
|
|||
|
|
@ -24,3 +24,28 @@ self.addEventListener('fetch', (event) => {
|
|||
})
|
||||
);
|
||||
});
|
||||
self.addEventListener('push', (event) => {
|
||||
const data = event.data ? event.data.json() : { title: 'Notification', body: 'New update from Agency' };
|
||||
|
||||
const options = {
|
||||
body: data.body,
|
||||
icon: '/icon-192.png',
|
||||
badge: '/icon-192.png',
|
||||
vibrate: [100, 50, 100],
|
||||
data: {
|
||||
dateOfArrival: Date.now(),
|
||||
primaryKey: '1'
|
||||
}
|
||||
};
|
||||
|
||||
event.waitUntil(
|
||||
self.registration.showNotification(data.title, options)
|
||||
);
|
||||
});
|
||||
|
||||
self.addEventListener('notificationclick', (event) => {
|
||||
event.notification.close();
|
||||
event.waitUntil(
|
||||
clients.openWindow('/')
|
||||
);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ mod m20220101_000001_create_table;
|
|||
mod m20260210_000002_add_answer_column;
|
||||
mod m20260210_000003_separate_runs;
|
||||
mod m20260210_000004_add_cron_column;
|
||||
mod m20260212_000005_notifications;
|
||||
|
||||
pub struct Migrator;
|
||||
|
||||
|
|
@ -15,6 +16,7 @@ impl MigratorTrait for Migrator {
|
|||
Box::new(m20260210_000002_add_answer_column::Migration),
|
||||
Box::new(m20260210_000003_separate_runs::Migration),
|
||||
Box::new(m20260210_000004_add_cron_column::Migration),
|
||||
Box::new(m20260212_000005_notifications::Migration),
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
115
migration/src/m20260212_000005_notifications.rs
Normal file
115
migration/src/m20260212_000005_notifications.rs
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
use sea_orm_migration::prelude::*;
|
||||
|
||||
#[derive(DeriveMigrationName)]
|
||||
pub struct Migration;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl MigrationTrait for Migration {
|
||||
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
|
||||
// Push Subscriptions table
|
||||
manager
|
||||
.create_table(
|
||||
Table::create()
|
||||
.table(PushSubscriptions::Table)
|
||||
.if_not_exists()
|
||||
.col(
|
||||
ColumnDef::new(PushSubscriptions::Id)
|
||||
.uuid()
|
||||
.not_null()
|
||||
.primary_key(),
|
||||
)
|
||||
.col(
|
||||
ColumnDef::new(PushSubscriptions::UserSub)
|
||||
.string()
|
||||
.not_null(),
|
||||
)
|
||||
.col(
|
||||
ColumnDef::new(PushSubscriptions::Endpoint)
|
||||
.string()
|
||||
.not_null(),
|
||||
)
|
||||
.col(
|
||||
ColumnDef::new(PushSubscriptions::P256dh)
|
||||
.string()
|
||||
.not_null(),
|
||||
)
|
||||
.col(ColumnDef::new(PushSubscriptions::Auth).string().not_null())
|
||||
.col(
|
||||
ColumnDef::new(PushSubscriptions::CreatedAt)
|
||||
.timestamp_with_time_zone()
|
||||
.not_null(),
|
||||
)
|
||||
.to_owned(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Task Subscriptions table
|
||||
manager
|
||||
.create_table(
|
||||
Table::create()
|
||||
.table(TaskSubscriptions::Table)
|
||||
.if_not_exists()
|
||||
.col(
|
||||
ColumnDef::new(TaskSubscriptions::Id)
|
||||
.uuid()
|
||||
.not_null()
|
||||
.primary_key(),
|
||||
)
|
||||
.col(
|
||||
ColumnDef::new(TaskSubscriptions::UserSub)
|
||||
.string()
|
||||
.not_null(),
|
||||
)
|
||||
.col(ColumnDef::new(TaskSubscriptions::TaskId).uuid().not_null())
|
||||
.col(
|
||||
ColumnDef::new(TaskSubscriptions::CreatedAt)
|
||||
.timestamp_with_time_zone()
|
||||
.not_null(),
|
||||
)
|
||||
.foreign_key(
|
||||
ForeignKey::create()
|
||||
.name("fk-task-subscription-task-id")
|
||||
.from(TaskSubscriptions::Table, TaskSubscriptions::TaskId)
|
||||
.to(Tasks::Table, Tasks::Id)
|
||||
.on_delete(ForeignKeyAction::Cascade),
|
||||
)
|
||||
.to_owned(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
|
||||
manager
|
||||
.drop_table(Table::drop().table(TaskSubscriptions::Table).to_owned())
|
||||
.await?;
|
||||
manager
|
||||
.drop_table(Table::drop().table(PushSubscriptions::Table).to_owned())
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(DeriveIden)]
|
||||
enum PushSubscriptions {
|
||||
Table,
|
||||
Id,
|
||||
UserSub,
|
||||
Endpoint,
|
||||
P256dh,
|
||||
Auth,
|
||||
CreatedAt,
|
||||
}
|
||||
|
||||
#[derive(DeriveIden)]
|
||||
enum TaskSubscriptions {
|
||||
Table,
|
||||
Id,
|
||||
UserSub,
|
||||
TaskId,
|
||||
CreatedAt,
|
||||
}
|
||||
|
||||
#[derive(DeriveIden)]
|
||||
enum Tasks {
|
||||
Table,
|
||||
Id,
|
||||
}
|
||||
|
|
@ -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