Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions src/apps/desktop/src/api/insights_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use serde::Deserialize;
#[serde(rename_all = "camelCase")]
pub struct GenerateInsightsRequest {
pub days: Option<u32>,
pub model_id: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
Expand All @@ -17,12 +18,18 @@ pub struct LoadInsightsReportRequest {
#[tauri::command]
pub async fn generate_insights(request: GenerateInsightsRequest) -> Result<InsightsReport, String> {
let days = request.days.unwrap_or(30);
info!("Generating insights for the last {} days", days);
let model_id = request.model_id;
info!(
"Generating insights for the last {} days with model selector {:?}",
days, model_id
);

InsightsService::generate(days).await.map_err(|e| {
error!("Failed to generate insights: {}", e);
format!("Failed to generate insights: {}", e)
})
InsightsService::generate(days, model_id)
.await
.map_err(|e| {
error!("Failed to generate insights: {}", e);
format!("Failed to generate insights: {}", e)
})
}

#[tauri::command]
Expand Down
16 changes: 5 additions & 11 deletions src/apps/desktop/src/api/remote_workspace_policy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] =
("cancel_dialog_turn", RemoteWorkspacePolicy::LegacyUnaudited),
(
"cancel_insights_generation",
RemoteWorkspacePolicy::LegacyUnaudited,
RemoteWorkspacePolicy::LocalOnly,
),
(
"cancel_mcp_remote_oauth",
Expand Down Expand Up @@ -383,7 +383,7 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] =
"generate_greeting_only",
RemoteWorkspacePolicy::LegacyUnaudited,
),
("generate_insights", RemoteWorkspacePolicy::LegacyUnaudited),
("generate_insights", RemoteWorkspacePolicy::RemoteRouted),
(
"generate_session_title",
RemoteWorkspacePolicy::LegacyUnaudited,
Expand Down Expand Up @@ -481,10 +481,7 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] =
"get_health_status",
RemoteWorkspacePolicy::WorkspaceAgnostic,
),
(
"get_latest_insights",
RemoteWorkspacePolicy::LegacyUnaudited,
),
("get_latest_insights", RemoteWorkspacePolicy::LocalOnly),
("get_mcp_prompt", RemoteWorkspacePolicy::LegacyUnaudited),
(
"get_mcp_remote_oauth_session",
Expand Down Expand Up @@ -648,7 +645,7 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] =
"grant_miniapp_workspace",
RemoteWorkspacePolicy::LegacyUnaudited,
),
("has_insights_data", RemoteWorkspacePolicy::LegacyUnaudited),
("has_insights_data", RemoteWorkspacePolicy::RemoteRouted),
(
"hide_agent_companion_desktop_pet",
RemoteWorkspacePolicy::LocalOnly,
Expand Down Expand Up @@ -770,10 +767,7 @@ pub const REMOTE_WORKSPACE_COMMAND_POLICIES: &[(&str, RemoteWorkspacePolicy)] =
"load_git_repo_history",
RemoteWorkspacePolicy::LegacyUnaudited,
),
(
"load_insights_report",
RemoteWorkspacePolicy::LegacyUnaudited,
),
("load_insights_report", RemoteWorkspacePolicy::LocalOnly),
(
"load_mcp_json_config",
RemoteWorkspacePolicy::LegacyUnaudited,
Expand Down
52 changes: 44 additions & 8 deletions src/crates/assembly/core/src/agentic/insights/cancellation.rs
Original file line number Diff line number Diff line change
@@ -1,27 +1,39 @@
use log::{debug, info, warn};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;

type Slot = Arc<Mutex<Option<CancellationToken>>>;
#[derive(Clone)]
pub struct Registration {
pub id: u64,
pub token: CancellationToken,
}

type Slot = Arc<Mutex<Option<Registration>>>;

static SLOT: std::sync::OnceLock<Slot> = std::sync::OnceLock::new();
static NEXT_ID: AtomicU64 = AtomicU64::new(1);

fn get_slot() -> Slot {
SLOT.get_or_init(|| Arc::new(Mutex::new(None))).clone()
}

/// Registers a new insights generation task, cancelling any previous one.
pub async fn register() -> CancellationToken {
pub async fn register() -> Registration {
let token = CancellationToken::new();
let registration = Registration {
id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
token: token.clone(),
};
let arc = get_slot();
let mut slot = arc.lock().await;
if let Some(old) = slot.take() {
old.cancel();
old.token.cancel();
debug!("Cancelled previous insights generation");
}
*slot = Some(token.clone());
token
*slot = Some(registration.clone());
registration
}

/// Cancels the current insights generation task.
Expand All @@ -30,7 +42,7 @@ pub async fn cancel() -> Result<(), String> {
let mut slot = arc.lock().await;
match slot.take() {
Some(token) => {
token.cancel();
token.token.cancel();
info!("Insights generation cancelled by user");
Ok(())
}
Expand All @@ -42,8 +54,32 @@ pub async fn cancel() -> Result<(), String> {
}

/// Unregisters the current task (call on completion).
pub async fn unregister() {
pub async fn unregister(registration_id: u64) {
let arc = get_slot();
let mut slot = arc.lock().await;
*slot = None;
if slot
.as_ref()
.is_some_and(|current| current.id == registration_id)
{
*slot = None;
}
}

#[cfg(test)]
mod tests {
use super::*;

#[tokio::test]
async fn old_task_cannot_unregister_a_new_generation() {
let old = register().await;
let current = register().await;
assert!(old.token.is_cancelled());

unregister(old.id).await;
cancel()
.await
.expect("current generation remains registered");

assert!(current.token.is_cancelled());
}
}
Loading
Loading