Reset prevent-sleep cap on agent activity (#1335)

Signed-off-by: Wes <wesbillman@users.noreply.github.com>
Co-authored-by: Pinky <44b8e82baa6e0e254e0208d68f335c283c94e7b78dd1fa10d5a49d3f13dd0435@sprout-oss.stage.blox.sqprod.co>

desktop/src-tauri/src/prevent_sleep.rs

@@ -8,6 +8,7 @@ use tauri::{AppHandle, Emitter};
 pub struct PreventSleepState {
     assertion_id: Option<u32>,
     timer_handle: Option<tauri::async_runtime::JoinHandle<()>>,
+    timer_generation: u64,
 }
 
 // ── macOS implementation ────────────────────────────────────────────────────
@@ -43,19 +44,68 @@ const K_IOPM_ASSERTION_LEVEL_ON: u32 = 255;
 #[cfg(target_os = "macos")]
 const K_CF_STRING_ENCODING_UTF8: u32 = 0x0800_0100;
 
-/// 4-hour cap in seconds.
-const CAP_SECONDS: u64 = 4 * 3600;
+/// One hour covers a silent long-running tool call while bounding idle keep-awake time.
+const INACTIVITY_CAP_SECONDS: u64 = 60 * 60;
+
+fn arm_cap_timer(
+    guard: &mut PreventSleepState,
+    state: &Arc<Mutex<PreventSleepState>>,
+    app_handle: &AppHandle,
+) {
+    if let Some(handle) = guard.timer_handle.take() {
+        handle.abort();
+    }
+
+    guard.timer_generation = guard.timer_generation.wrapping_add(1);
+    let generation = guard.timer_generation;
+    let handle = app_handle.clone();
+    let timer_state = Arc::clone(state);
+    let timer_task = tauri::async_runtime::spawn(async move {
+        tokio::time::sleep(std::time::Duration::from_secs(INACTIVITY_CAP_SECONDS)).await;
+        if expire_if_current(&timer_state, generation) {
+            let _ = handle.emit("prevent-sleep-expired", ());
+        }
+    });
+    guard.timer_handle = Some(timer_task);
+}
+
+fn expire_if_current(state: &Arc<Mutex<PreventSleepState>>, generation: u64) -> bool {
+    let mut guard = match state.lock() {
+        Ok(g) => g,
+        Err(_) => return false,
+    };
+
+    if guard.timer_generation != generation {
+        return false;
+    }
+
+    guard.timer_handle = None;
+
+    #[cfg(target_os = "macos")]
+    if let Some(id) = guard.assertion_id.take() {
+        unsafe {
+            macos::IOPMAssertionRelease(id);
+        }
+    }
+
+    #[cfg(not(target_os = "macos"))]
+    {
+        guard.assertion_id = None;
+    }
+
+    true
+}
 
 /// Create a `PreventUserIdleSystemSleep` assertion if not already held.
-/// Starts a 4-hour timer that auto-releases and emits `prevent-sleep-expired`.
+/// Refreshes the inactivity cap when the assertion is already held.
 pub fn acquire(
     state: &Arc<Mutex<PreventSleepState>>,
     app_handle: &AppHandle,
 ) -> Result<(), String> {
     let mut guard = state.lock().map_err(|e| e.to_string())?;
 
-    // Idempotent — already held.
     if guard.assertion_id.is_some() {
+        arm_cap_timer(&mut guard, state, app_handle);
         return Ok(());
     }
 
@@ -107,16 +157,9 @@ pub fn acquire(
         }
     }
 
-    // Start the 4-hour cap timer only if an assertion was actually created.
+    // Start the inactivity cap timer only if an assertion was actually created.
     if guard.assertion_id.is_some() {
-        let handle = app_handle.clone();
-        let timer_state = Arc::clone(state);
-        let timer_task = tauri::async_runtime::spawn(async move {
-            tokio::time::sleep(std::time::Duration::from_secs(CAP_SECONDS)).await;
-            release(&timer_state);
-            let _ = handle.emit("prevent-sleep-expired", ());
-        });
-        guard.timer_handle = Some(timer_task);
+        arm_cap_timer(&mut guard, state, app_handle);
     }
 
     Ok(())
@@ -132,6 +175,7 @@ pub fn release(state: &Arc<Mutex<PreventSleepState>>) {
     if let Some(handle) = guard.timer_handle.take() {
         handle.abort();
     }
+    guard.timer_generation = guard.timer_generation.wrapping_add(1);
 
     #[cfg(target_os = "macos")]
     if let Some(id) = guard.assertion_id.take() {

desktop/src/features/agents/preventSleepActivity.test.mjs

@@ -0,0 +1,37 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import { createPreventSleepActivityTracker } from "./preventSleepActivity.ts";
+
+const AGENT =
+  "abcd1234abcd1234abcd1234abcd1234abcd1234abcd1234abcd1234abcd1234";
+
+function event(seq) {
+  return { seq, timestamp: `2024-01-01T00:00:0${seq}Z` };
+}
+
+test("prevent sleep activity tracker seeds existing observer events without touching", () => {
+  const tracker = createPreventSleepActivityTracker();
+
+  assert.equal(tracker.observe([{ pubkey: AGENT, events: [event(1)] }]), false);
+});
+
+test("prevent sleep activity tracker reports a newer observer event", () => {
+  const tracker = createPreventSleepActivityTracker();
+  tracker.observe([{ pubkey: AGENT, events: [event(1)] }]);
+
+  assert.equal(
+    tracker.observe([{ pubkey: AGENT, events: [event(1), event(2)] }]),
+    true,
+  );
+});
+
+test("prevent sleep activity tracker ignores unchanged latest observer event", () => {
+  const tracker = createPreventSleepActivityTracker();
+  tracker.observe([{ pubkey: AGENT, events: [event(1), event(2)] }]);
+
+  assert.equal(
+    tracker.observe([{ pubkey: AGENT, events: [event(1), event(2)] }]),
+    false,
+  );
+});

desktop/src/features/agents/preventSleepActivity.ts

@@ -0,0 +1,40 @@
+type ObserverEventSummary = {
+  seq: number;
+  timestamp: string;
+};
+
+type AgentObserverEvents = {
+  pubkey: string;
+  events: readonly ObserverEventSummary[];
+};
+
+function latestEventKey(
+  events: readonly ObserverEventSummary[],
+): string | null {
+  const latest = events.length > 0 ? events[events.length - 1] : undefined;
+  return latest ? `${latest.timestamp}:${latest.seq}` : null;
+}
+
+export function createPreventSleepActivityTracker() {
+  const lastEventByAgent = new Map<string, string>();
+
+  return {
+    observe(agents: readonly AgentObserverEvents[]): boolean {
+      let hasNewActivity = false;
+
+      for (const agent of agents) {
+        const latestKey = latestEventKey(agent.events);
+        if (!latestKey) continue;
+
+        const previousKey = lastEventByAgent.get(agent.pubkey);
+        lastEventByAgent.set(agent.pubkey, latestKey);
+
+        if (previousKey !== undefined && previousKey !== latestKey) {
+          hasNewActivity = true;
+        }
+      }
+
+      return hasNewActivity;
+    },
+  };
+}

desktop/src/features/agents/usePreventSleep.ts

@@ -1,6 +1,13 @@
 import * as React from "react";
 import { useManagedAgentsQuery } from "@/features/agents/hooks";
+import {
+  getAgentObserverSnapshot,
+  subscribeAgentObserverStore,
+  useManagedAgentObserverBridge,
+} from "@/features/agents/observerRelayStore";
+import { createPreventSleepActivityTracker } from "@/features/agents/preventSleepActivity";
 import { setPreventSleepActive } from "@/shared/api/tauri";
+import { normalizePubkey } from "@/shared/lib/pubkey";
 import { listen } from "@tauri-apps/api/event";
 
 // Intentionally not scoped per-pubkey — multi-user desktop is rare and the
@@ -51,11 +58,31 @@ function usePreventSleepInternal() {
 
   // Only local "running" agents need sleep prevention. Remote "deployed"
   // agents run on provider infrastructure and are unaffected by local sleep.
-  const hasRunningAgents = React.useMemo(
-    () => agents?.some((agent) => agent.status === "running") ?? false,
+  const runningAgentPubkeys = React.useMemo(
+    () =>
+      (agents ?? [])
+        .filter((agent) => agent.status === "running")
+        .map((agent) => normalizePubkey(agent.pubkey))
+        .sort(),
     [agents],
   );
 
+  const runningAgentPubkeyKey = runningAgentPubkeys.join(",");
+  const runningObserverAgents = React.useMemo(
+    () =>
+      enabled && runningAgentPubkeyKey
+        ? runningAgentPubkeyKey.split(",").map((pubkey) => ({
+            pubkey,
+            status: "running" as const,
+          }))
+        : [],
+    [enabled, runningAgentPubkeyKey],
+  );
+
+  useManagedAgentObserverBridge(runningObserverAgents);
+
+  const hasRunningAgents = runningAgentPubkeys.length > 0;
+
   const [expired, setExpired] = React.useState(false);
 
   const active = enabled && hasRunningAgents && !expired;
@@ -77,6 +104,30 @@ function usePreventSleepInternal() {
     };
   }, []);
 
+  React.useEffect(() => {
+    if (!enabled || !runningAgentPubkeyKey) return;
+
+    const observedPubkeys = runningAgentPubkeyKey.split(",");
+    const tracker = createPreventSleepActivityTracker();
+    const observeActivity = () => {
+      const hasNewActivity = tracker.observe(
+        observedPubkeys.map((pubkey) => ({
+          pubkey,
+          events: getAgentObserverSnapshot(pubkey, true).events,
+        })),
+      );
+      if (!hasNewActivity) return;
+
+      if (expired) {
+        setExpired(false);
+      }
+      void setPreventSleepActive(true);
+    };
+
+    observeActivity();
+    return subscribeAgentObserverStore(observeActivity);
+  }, [enabled, expired, runningAgentPubkeyKey]);
+
   return {
     enabled,
     setEnabled,

desktop/src/features/settings/ui/PreventSleepSettingsCard.tsx

@@ -25,8 +25,8 @@ export function PreventSleepSettingsCard() {
             </label>
             <p className="text-sm font-normal text-muted-foreground">
               Prevents your computer from sleeping while local agents are
-              running. Automatically releases when all agents stop or after 4
-              hours.
+              running. Automatically releases when all agents stop or after 1
+              hour without agent activity.
             </p>
           </div>
           <Switch
@@ -51,8 +51,9 @@ export function PreventSleepSettingsCard() {
 
       {expired && (
         <p className="mt-3 rounded-xl border border-yellow-500/30 bg-yellow-500/10 px-3 py-2 text-sm text-yellow-700 dark:text-yellow-400">
-          Sleep prevention expired after 4 hours. Toggle off and on to
-          re-enable.
+          Sleep prevention expired after 1 hour without agent activity. It will
+          resume on the next agent activity, or toggle off and on to re-enable
+          now.
         </p>
       )}
     </section>