diff --git a/src/openhuman/subconscious/heartbeat/planner/collectors.rs b/src/openhuman/subconscious/heartbeat/planner/collectors.rs index a1e5d4657..187ede703 100644 --- a/src/openhuman/subconscious/heartbeat/planner/collectors.rs +++ b/src/openhuman/subconscious/heartbeat/planner/collectors.rs @@ -1,4 +1,6 @@ +use async_trait::async_trait; use chrono::{DateTime, Duration, Utc}; +use futures::stream::StreamExt; use serde_json::json; use crate::openhuman::composio::client::{ @@ -148,6 +150,180 @@ fn select_calendar_connections_for_tick( selected } +/// Bound on how many calendar connections are polled concurrently per tick. +/// +/// `select_calendar_connections_for_tick` already caps the selected set to +/// `max_calendar_connections_per_tick` (default 2), so this only matters when a +/// user raises that cap. Kept small to avoid opening more than a handful of +/// sockets against Composio/Google Calendar in a single heartbeat tick. +const CALENDAR_FANOUT_CONCURRENCY: usize = 4; + +/// Narrow seam over the mode-aware Composio client so the calendar fan-out can +/// be unit-tested with a fake executor (no network). Both real variants +/// (`Backend` / `Direct`) and the fake forward through the identical +/// arg-building + extraction path. +#[async_trait] +trait CalendarExecutor { + /// Stable label for the underlying client variant, surfaced on the failure + /// log so backend vs direct dispatch stays distinguishable post-refactor. + fn kind_label(&self) -> &'static str; + + /// Execute one `GOOGLECALENDAR_EVENTS_LIST` call for a single connection. + async fn execute( + &self, + slug: &str, + arguments: Option, + ) -> anyhow::Result; +} + +/// Real executor wrapping the resolved mode-aware client. Holds borrows only — +/// constructed per tick, dropped when the fan-out completes. +struct ComposioCalendarExecutor<'a> { + kind: &'a ComposioClientKind, + entity_id: &'a str, +} + +#[async_trait] +impl CalendarExecutor for ComposioCalendarExecutor<'_> { + fn kind_label(&self) -> &'static str { + match self.kind { + ComposioClientKind::Backend(_) => "backend", + ComposioClientKind::Direct(_) => "direct", + } + } + + async fn execute( + &self, + slug: &str, + arguments: Option, + ) -> anyhow::Result { + match self.kind { + ComposioClientKind::Backend(client) => client.execute_tool(slug, arguments).await, + ComposioClientKind::Direct(direct) => { + direct_execute(direct, slug, arguments, self.entity_id, None).await + } + } + } +} + +/// Fetch + extract upcoming meetings for one calendar connection. +/// +/// A failed fetch contributes no events (returns an empty `Vec`), matching the +/// serial loop's `continue` so one broken connection never poisons the tick. +async fn fetch_calendar_events_for_connection( + executor: &(dyn CalendarExecutor + Sync), + conn: &ComposioConnection, + meeting_lookahead_minutes: u32, + now: DateTime, + end_window: DateTime, +) -> Vec { + let toolkit = conn.normalized_toolkit(); + + // Build base args, then let the shared transformer fill in `timeZone` + + // `singleEvents` so this poller behaves identically to the agent-driven + // dispatcher path (issue #1714). Routing both call sites through the same + // helper means a future change to the defaulting policy only has to land in + // one place. + let arguments = json!({ + "connectionId": conn.id, + "timeMin": now.to_rfc3339(), + "timeMax": end_window.to_rfc3339(), + "maxResults": 20 + }); + let iana = crate::openhuman::composio::googlecalendar_args::current_iana_timezone(); + tracing::debug!( + target: "composio", + slug = "GOOGLECALENDAR_EVENTS_LIST", + toolkit = %toolkit, + connection_id = %conn.id, + iana = %iana, + lookahead_minutes = meeting_lookahead_minutes, + "[composio][heartbeat-planner] applying calendar query defaults pre-poll" + ); + let arguments = crate::openhuman::composio::googlecalendar_args::apply_calendar_query_defaults( + "GOOGLECALENDAR_EVENTS_LIST", + Some(arguments), + &iana, + ); + + match executor + .execute("GOOGLECALENDAR_EVENTS_LIST", arguments) + .await + { + // Composio encodes provider-side failures in `successful`/`error` while + // still returning `Ok(_)` (the repo-wide `if !resp.successful` convention), + // so an unsuccessful list must take the same warn + empty path as a + // transport `Err` — otherwise an error payload falls through to + // `extract_calendar_events`, losing the diagnostic and risking stale data. + Ok(resp) if resp.successful => { + extract_calendar_events(&resp.data, &toolkit, &conn.id, now, end_window) + } + Ok(resp) => { + tracing::warn!( + target: "composio", + toolkit = %toolkit, + connection_id = %conn.id, + kind = executor.kind_label(), + error = %resp + .error + .as_deref() + .unwrap_or("calendar execute returned unsuccessful=false"), + "[heartbeat:planner] GOOGLECALENDAR_EVENTS_LIST failed" + ); + Vec::new() + } + Err(error) => { + tracing::warn!( + target: "composio", + toolkit = %toolkit, + connection_id = %conn.id, + kind = executor.kind_label(), + error = %error, + "[heartbeat:planner] GOOGLECALENDAR_EVENTS_LIST failed" + ); + Vec::new() + } + } +} + +/// Drive the per-connection fetches with bounded concurrency. +/// +/// `buffered(K)` yields results in input order, so the flattened event stream is +/// identical to the old strictly-serial loop; only the wall-clock latency drops +/// (from the sum of all fetches to roughly `ceil(N / K)` round-trips). Single +/// connections (the common case) cost nothing extra. +async fn collect_calendar_events_buffered( + executor: &(dyn CalendarExecutor + Sync), + connections: &[ComposioConnection], + meeting_lookahead_minutes: u32, + now: DateTime, + end_window: DateTime, +) -> Vec { + // Materialize the per-connection futures into a `Vec` before `stream::iter` + // — handing it a lazy `Map` adaptor trips the "implementation of `Send` is + // not general enough" HRTB error. + let fetch_futs: Vec<_> = connections + .iter() + .map(|conn| { + fetch_calendar_events_for_connection( + executor, + conn, + meeting_lookahead_minutes, + now, + end_window, + ) + }) + .collect(); + + futures::stream::iter(fetch_futs) + .buffered(CALENDAR_FANOUT_CONCURRENCY) + .collect::>>() + .await + .into_iter() + .flatten() + .collect() +} + pub(crate) async fn collect_calendar_meetings( config: &Config, now: DateTime, @@ -208,96 +384,68 @@ pub(crate) async fn collect_calendar_meetings( }, }; - let lookahead = Duration::minutes(i64::from(config.heartbeat.meeting_lookahead_minutes.max(1))); - let end_window = now + lookahead; + let executor = ComposioCalendarExecutor { + kind: &kind, + entity_id: &config.composio.entity_id, + }; + collect_calendar_meetings_with( + &executor, + connections, + now, + config.heartbeat.meeting_lookahead_minutes.max(1), + config.heartbeat.max_calendar_connections_per_tick.max(1) as usize, + config.heartbeat.interval_minutes, + ) + .await +} - let mut out = Vec::new(); - let calendar_connection_limit = - config.heartbeat.max_calendar_connections_per_tick.max(1) as usize; - for conn in select_calendar_connections_for_tick( +/// Tick-rotation selection + bounded fan-out, decoupled from the concrete +/// Composio client so it is unit-testable with a fake `CalendarExecutor`. +/// +/// `collect_calendar_meetings` only builds the real executor and delegates here; +/// the rotation cap (`select_calendar_connections_for_tick`) and lookahead window +/// run exactly as before. +async fn collect_calendar_meetings_with( + executor: &(dyn CalendarExecutor + Sync), + connections: Vec, + now: DateTime, + meeting_lookahead_minutes: u32, + calendar_connection_limit: usize, + interval_minutes: u32, +) -> Vec { + let end_window = now + Duration::minutes(i64::from(meeting_lookahead_minutes)); + let selected = select_calendar_connections_for_tick( connections, calendar_connection_limit, now, - config.heartbeat.interval_minutes, - ) { - let toolkit = conn.normalized_toolkit(); - - // Build base args, then let the shared transformer fill in - // `timeZone` + `singleEvents` so this poller behaves identically - // to the agent-driven dispatcher path (issue #1714). Routing - // both call sites through the same helper means a future change - // to the defaulting policy only has to land in one place. - let arguments = json!({ - "connectionId": conn.id, - "timeMin": now.to_rfc3339(), - "timeMax": end_window.to_rfc3339(), - "maxResults": 20 - }); - let iana = crate::openhuman::composio::googlecalendar_args::current_iana_timezone(); - tracing::debug!( - target: "composio", - slug = "GOOGLECALENDAR_EVENTS_LIST", - toolkit = %toolkit, - connection_id = %conn.id, - iana = %iana, - lookahead_minutes = config.heartbeat.meeting_lookahead_minutes.max(1), - "[composio][heartbeat-planner] applying calendar query defaults pre-poll" - ); - let arguments = - crate::openhuman::composio::googlecalendar_args::apply_calendar_query_defaults( - "GOOGLECALENDAR_EVENTS_LIST", - Some(arguments), - &iana, - ); - - let resp: ComposioExecuteResponse = match &kind { - ComposioClientKind::Backend(client) => { - match client - .execute_tool("GOOGLECALENDAR_EVENTS_LIST", arguments) - .await - { - Ok(resp) => resp, - Err(error) => { - tracing::warn!( - toolkit = %toolkit, - connection_id = %conn.id, - error = %error, - "[heartbeat:planner] GOOGLECALENDAR_EVENTS_LIST (backend) failed" - ); - continue; - } - } - } - ComposioClientKind::Direct(direct) => { - match direct_execute( - direct, - "GOOGLECALENDAR_EVENTS_LIST", - arguments, - &config.composio.entity_id, - None, - ) - .await - { - Ok(resp) => resp, - Err(error) => { - tracing::warn!( - toolkit = %toolkit, - connection_id = %conn.id, - error = %error, - "[heartbeat:planner] GOOGLECALENDAR_EVENTS_LIST (direct) failed" - ); - continue; - } - } - } - }; - - out.extend(extract_calendar_events( - &resp.data, &toolkit, &conn.id, now, end_window, - )); - } - - out + interval_minutes, + ); + tracing::debug!( + executor = executor.kind_label(), + selected = selected.len(), + calendar_connection_limit, + meeting_lookahead_minutes, + interval_minutes, + concurrency = CALENDAR_FANOUT_CONCURRENCY, + now = %now, + end_window = %end_window, + "[heartbeat:planner] calendar fan-out start" + ); + let events = collect_calendar_events_buffered( + executor, + &selected, + meeting_lookahead_minutes, + now, + end_window, + ) + .await; + tracing::debug!( + executor = executor.kind_label(), + selected = selected.len(), + events = events.len(), + "[heartbeat:planner] calendar fan-out complete" + ); + events } pub(crate) fn extract_calendar_events( @@ -647,6 +795,257 @@ mod tests { assert_eq!(selected, vec!["active-cal"]); } + // ── Calendar fan-out (bounded concurrency) ─────────────────────────── + + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + /// Build a Composio `data` payload carrying a single in-window event whose + /// summary becomes the extracted `PendingEvent` title. + fn in_window_event(summary: &str, start: DateTime) -> serde_json::Value { + json!({ + "items": [ + { + "id": summary, + "summary": summary, + "start": { "dateTime": start.to_rfc3339() } + } + ] + }) + } + + /// In-memory `CalendarExecutor` keyed by `connectionId`, with an optional + /// per-call delay and an observed-max-concurrency probe. + struct FakeExecutor { + responses: std::collections::HashMap>, + /// Connections whose execute returns `Ok(successful=false)` carrying the + /// given would-be-event payload — models a provider-side failure that + /// still arrives as `Ok(_)`. + unsuccessful: std::collections::HashMap, + inflight: Arc, + max_inflight: Arc, + delay_ms: u64, + } + + impl FakeExecutor { + fn new(delay_ms: u64) -> Self { + Self { + responses: std::collections::HashMap::new(), + unsuccessful: std::collections::HashMap::new(), + inflight: Arc::new(AtomicUsize::new(0)), + max_inflight: Arc::new(AtomicUsize::new(0)), + delay_ms, + } + } + + fn with(mut self, conn_id: &str, resp: Result) -> Self { + self.responses.insert(conn_id.to_string(), resp); + self + } + + /// Register a connection that returns `Ok` but with `successful=false`, + /// carrying `data` that *would* parse to an event if it were extracted. + fn with_unsuccessful(mut self, conn_id: &str, data: serde_json::Value) -> Self { + self.unsuccessful.insert(conn_id.to_string(), data); + self + } + } + + #[async_trait] + impl CalendarExecutor for FakeExecutor { + fn kind_label(&self) -> &'static str { + "fake" + } + + async fn execute( + &self, + _slug: &str, + arguments: Option, + ) -> anyhow::Result { + let current = self.inflight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_inflight.fetch_max(current, Ordering::SeqCst); + if self.delay_ms > 0 { + tokio::time::sleep(std::time::Duration::from_millis(self.delay_ms)).await; + } + self.inflight.fetch_sub(1, Ordering::SeqCst); + + let conn_id = arguments + .as_ref() + .and_then(|a| a.get("connectionId")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + if let Some(data) = self.unsuccessful.get(&conn_id) { + return Ok(ComposioExecuteResponse { + data: data.clone(), + successful: false, + error: Some("provider rejected the request".to_string()), + cost_usd: 0.0, + markdown_formatted: None, + }); + } + match self.responses.get(&conn_id) { + Some(Ok(data)) => Ok(ComposioExecuteResponse { + data: data.clone(), + successful: true, + error: None, + cost_usd: 0.0, + markdown_formatted: None, + }), + Some(Err(msg)) => Err(anyhow::anyhow!(msg.clone())), + None => Err(anyhow::anyhow!("no canned response for {conn_id}")), + } + } + } + + #[tokio::test] + async fn calendar_fanout_preserves_connection_order() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let end_window = now + Duration::minutes(120); + let start = now + Duration::minutes(30); + + // 5 connections > CALENDAR_FANOUT_CONCURRENCY (4), so the window wraps. + let conns: Vec = (1..=5) + .map(|n| conn(&format!("cal-{n}"), "googlecalendar", "ACTIVE")) + .collect(); + + let mut exec = FakeExecutor::new(0); + for n in 1..=5 { + exec = exec.with( + &format!("cal-{n}"), + Ok(in_window_event(&format!("evt-{n}"), start)), + ); + } + + let events = collect_calendar_events_buffered(&exec, &conns, 120, now, end_window).await; + let titles: Vec = events.into_iter().map(|e| e.title).collect(); + assert_eq!( + titles, + vec!["evt-1", "evt-2", "evt-3", "evt-4", "evt-5"], + "buffered fan-out must preserve per-connection input order" + ); + } + + #[tokio::test] + async fn calendar_fanout_runs_connections_concurrently() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let end_window = now + Duration::minutes(120); + let start = now + Duration::minutes(30); + + // Drive *more* connections than the cap (N = CALENDAR_FANOUT_CONCURRENCY + // + 2) so this test fails if the implementation ever regresses to an + // unbounded fan-out: with exactly `CALENDAR_FANOUT_CONCURRENCY` + // connections the cap is indistinguishable from no cap at all. + let total = CALENDAR_FANOUT_CONCURRENCY + 2; + let conns: Vec = (1..=total) + .map(|n| conn(&format!("cal-{n}"), "googlecalendar", "ACTIVE")) + .collect(); + + let mut exec = FakeExecutor::new(30); + for n in 1..=total { + exec = exec.with( + &format!("cal-{n}"), + Ok(in_window_event(&format!("evt-{n}"), start)), + ); + } + let max_inflight = exec.max_inflight.clone(); + + let events = collect_calendar_events_buffered(&exec, &conns, 120, now, end_window).await; + assert_eq!(events.len(), total); + let observed = max_inflight.load(Ordering::SeqCst); + assert!( + observed > 1, + "fan-out must overlap fetches (observed max in-flight = {observed})" + ); + assert!( + observed <= CALENDAR_FANOUT_CONCURRENCY, + "fan-out must stay within the concurrency cap of {CALENDAR_FANOUT_CONCURRENCY} \ + (observed max in-flight = {observed})" + ); + } + + #[tokio::test] + async fn calendar_fanout_isolates_failed_connection() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let end_window = now + Duration::minutes(120); + let start = now + Duration::minutes(30); + + let conns = vec![ + conn("cal-1", "googlecalendar", "ACTIVE"), + conn("cal-2", "googlecalendar", "ACTIVE"), + conn("cal-3", "googlecalendar", "ACTIVE"), + ]; + + // Middle connection errors — it must drop out without poisoning the + // surviving two, which keep their input order. + let exec = FakeExecutor::new(0) + .with("cal-1", Ok(in_window_event("evt-1", start))) + .with("cal-2", Err("boom".to_string())) + .with("cal-3", Ok(in_window_event("evt-3", start))); + + let events = collect_calendar_events_buffered(&exec, &conns, 120, now, end_window).await; + let titles: Vec = events.into_iter().map(|e| e.title).collect(); + assert_eq!( + titles, + vec!["evt-1", "evt-3"], + "a failed connection contributes no events while order is preserved" + ); + } + + #[tokio::test] + async fn calendar_fanout_drops_unsuccessful_response() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let end_window = now + Duration::minutes(120); + let start = now + Duration::minutes(30); + + let conns = vec![ + conn("cal-1", "googlecalendar", "ACTIVE"), + conn("cal-2", "googlecalendar", "ACTIVE"), + ]; + + // cal-2 returns Ok(successful=false) carrying a payload that *would* + // parse to an event — proving the unsuccessful path warns + drops it + // instead of falling through to `extract_calendar_events`. + let exec = FakeExecutor::new(0) + .with("cal-1", Ok(in_window_event("evt-1", start))) + .with_unsuccessful("cal-2", in_window_event("evt-2-must-not-appear", start)); + + let events = collect_calendar_events_buffered(&exec, &conns, 120, now, end_window).await; + let titles: Vec = events.into_iter().map(|e| e.title).collect(); + assert_eq!( + titles, + vec!["evt-1"], + "an Ok(successful=false) response must contribute no events, not its error payload" + ); + } + + #[tokio::test] + async fn collect_calendar_meetings_with_selects_and_fans_out() { + // ts=1_700_000_000 with interval=5 → tick_index=5_666_666, which is even, + // so offset = tick_index % 2 = 0: the two connections are selected in + // input order (identity rotation), and the fan-out preserves that order. + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let start = now + Duration::minutes(30); + + let conns = vec![ + conn("cal-1", "googlecalendar", "ACTIVE"), + conn("cal-2", "googlecalendar", "ACTIVE"), + ]; + + let exec = FakeExecutor::new(0) + .with("cal-1", Ok(in_window_event("evt-1", start))) + .with("cal-2", Ok(in_window_event("evt-2", start))); + + // limit=2 >= N so selection keeps both; interval=5. + let events = collect_calendar_meetings_with(&exec, conns, now, 120, 2, 5).await; + let titles: Vec = events.into_iter().map(|e| e.title).collect(); + assert_eq!( + titles, + vec!["evt-1", "evt-2"], + "wrapper must select the tick's connections and fan-out in order" + ); + } + // ── extract_meeting_url_from_map ───────────────────────────── fn map_from_value(v: serde_json::Value) -> serde_json::Map {