feat(memory-sync): track Composio API action count + cost in sync audit (#3111) (#3138)

Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>
This commit is contained in:
Justin
2026-06-01 20:55:44 -07:00
committed by GitHub
co-authored by Steven Enamakel
parent 5890dc56c7
commit 044728a7a4
21 changed files with 234 additions and 28 deletions
+2
View File
@@ -450,6 +450,7 @@ async fn main() -> Result<()> {
config: Arc::clone(&config),
toolkit: conn.toolkit.clone(),
connection_id: Some(conn.id.clone()),
usage: Default::default(),
};
match run_backfill_via_search(&ctx, cli.days).await {
Ok(outcome) => {
@@ -547,6 +548,7 @@ async fn main() -> Result<()> {
config: Arc::clone(&config),
toolkit: conn.toolkit.clone(),
connection_id: Some(conn.id.clone()),
usage: Default::default(),
};
match provider.sync(&ctx, SyncReason::Manual).await {
Ok(outcome) => {
+3
View File
@@ -1205,6 +1205,7 @@ pub async fn composio_get_user_profile(
config: Arc::new(config.clone()),
toolkit: toolkit.clone(),
connection_id: Some(connection_id.to_string()),
usage: Default::default(),
};
let profile = provider.fetch_user_profile(&ctx).await.map_err(|e| {
@@ -1276,6 +1277,7 @@ pub async fn composio_refresh_all_identities(
config: Arc::new(config.clone()),
toolkit: toolkit.clone(),
connection_id: Some(connection_id.clone()),
usage: Default::default(),
};
match provider.fetch_user_profile(&ctx).await {
@@ -1372,6 +1374,7 @@ pub async fn composio_sync(
config: Arc::new(config.clone()),
toolkit: toolkit.clone(),
connection_id: Some(connection_id.to_string()),
usage: Default::default(),
};
let started_at_ms = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
+7 -2
View File
@@ -168,7 +168,12 @@ async fn spawn_manual_sync(requested_connection: Option<String>) -> Result<(), S
)
.await
{
Ok(outcome) => {
// `run_connection_sync` returns `(SyncOutcome, ComposioUsage)`
// post-#3111; this caller only surfaces the outcome for UI
// stage events, so the usage tally is intentionally ignored
// here (the sync-audit caller in `memory_sources::sync` is the
// one that records it).
Ok((outcome, _usage)) => {
emit_sync_stage(
MemorySyncTrigger::Manual,
MemorySyncStage::Completed,
@@ -180,7 +185,7 @@ async fn spawn_manual_sync(requested_connection: Option<String>) -> Result<(), S
)),
);
}
Err(error) => {
Err((error, _usage)) => {
emit_sync_stage(
MemorySyncTrigger::Manual,
MemorySyncStage::Failed,
+26 -8
View File
@@ -24,7 +24,7 @@ use crate::openhuman::memory::sync::{emit_sync_stage, MemorySyncStage, MemorySyn
use crate::openhuman::memory_sources::readers;
use crate::openhuman::memory_sources::types::{MemorySourceEntry, SourceKind};
use crate::openhuman::memory_sync::canonicalize::document::DocumentInput;
use crate::openhuman::memory_sync::composio::{self, SyncReason};
use crate::openhuman::memory_sync::composio::{self, ComposioUsage, SyncReason};
const SYNC_CONCURRENCY: usize = 10;
@@ -89,8 +89,13 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<()
"[memory_sources:sync] dispatching by kind"
);
let sync_start = std::time::Instant::now();
// Composio billable-action usage for this run, populated by
// `sync_composio` (#3111). Stays zero for non-Composio kinds.
let mut composio_usage = ComposioUsage::default();
let outcome = match source.kind {
SourceKind::Composio => sync_composio(&source, config.clone()).await,
SourceKind::Composio => {
sync_composio(&source, config.clone(), &mut composio_usage).await
}
SourceKind::GithubRepo => {
// GitHub path writes its own detailed audit entry
// with token breakdowns; skip the dispatcher-level
@@ -150,6 +155,8 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<()
input_tokens: 0,
output_tokens: 0,
estimated_cost_usd: 0.0,
composio_actions_called: composio_usage.actions_called,
composio_cost_usd: composio_usage.cost_usd,
actual_charged_usd: None,
duration_ms,
success: true,
@@ -183,6 +190,8 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<()
input_tokens: 0,
output_tokens: 0,
estimated_cost_usd: 0.0,
composio_actions_called: composio_usage.actions_called,
composio_cost_usd: composio_usage.cost_usd,
actual_charged_usd: None,
duration_ms,
success: false,
@@ -226,7 +235,11 @@ pub async fn sync_source(source: MemorySourceEntry, config: Config) -> Result<()
Ok(())
}
async fn sync_composio(source: &MemorySourceEntry, config: Config) -> Result<usize, String> {
async fn sync_composio(
source: &MemorySourceEntry,
config: Config,
usage_out: &mut ComposioUsage,
) -> Result<usize, String> {
let connection_id = source
.connection_id
.as_deref()
@@ -240,11 +253,16 @@ async fn sync_composio(source: &MemorySourceEntry, config: Config) -> Result<usi
Some(format!("delegating to composio sync for {connection_id}")),
);
let outcome = composio::run_connection_sync(config, connection_id, SyncReason::Manual)
.await
.map_err(|e| format!("composio sync failed: {e}"))?;
Ok(outcome.items_ingested)
match composio::run_connection_sync(config, connection_id, SyncReason::Manual).await {
Ok((outcome, usage)) => {
*usage_out = usage;
Ok(outcome.items_ingested)
}
Err((e, usage)) => {
*usage_out = usage;
Err(format!("composio sync failed: {e}"))
}
}
}
/// Per-item sync path for Folder/RSS/WebPage sources.
+34 -8
View File
@@ -31,7 +31,7 @@ pub use periodic::{record_sync_success, start_periodic_sync};
pub use providers::{
all_providers as all_composio_sync_providers, get_provider as get_composio_sync_provider,
init_default_providers as init_default_composio_sync_providers, ComposioProvider,
ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason,
ComposioUsage, ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason,
};
/// One provider-backed connection that the memory sync layer can execute.
@@ -118,35 +118,61 @@ pub async fn scan_active_sync_targets(config: &Config) -> Result<Vec<SyncTarget>
}
/// Run one provider-backed sync end-to-end in-process.
///
/// Returns the provider's [`SyncOutcome`] together with the
/// [`ComposioUsage`] tally (billable action count + actual USD cost)
/// accumulated at the `execute` chokepoint during this run, so the
/// sync-audit caller can record Composio API-call cost alongside the LLM
/// summarisation cost (#3111).
pub async fn run_connection_sync(
config: Config,
connection_id: &str,
reason: SyncReason,
) -> Result<SyncOutcome, String> {
) -> Result<(SyncOutcome, ComposioUsage), (String, ComposioUsage)> {
init_default_composio_sync_providers();
let no_usage = |e: String| (e, ComposioUsage::default());
let target = list_sync_targets(&config)
.await?
.await
.map_err(no_usage)?
.into_iter()
.find(|target| target.connection_id == connection_id)
.ok_or_else(|| {
format!("no provider-backed active sync target for connection_id={connection_id}")
no_usage(format!(
"no provider-backed active sync target for connection_id={connection_id}",
))
})?;
let provider = get_composio_sync_provider(&target.toolkit).ok_or_else(|| {
format!(
no_usage(format!(
"no native memory sync provider registered for toolkit '{}'",
target.toolkit
)
target.toolkit,
))
})?;
let ctx = ProviderContext {
config: std::sync::Arc::new(config),
toolkit: target.toolkit,
connection_id: Some(target.connection_id),
usage: Default::default(),
};
provider.sync(&ctx, reason).await
let sync_result = provider.sync(&ctx, reason).await;
// Read the Composio billable-action tally *before* propagating errors.
// A sync that errors partway may still have fired billable actions;
// reading here ensures the dispatcher audit sees partial cost (#3111).
let usage = ctx
.usage
.lock()
.map(|u| u.clone())
.unwrap_or_else(|poisoned| poisoned.into_inner().clone());
match sync_result {
Ok(outcome) => Ok((outcome, usage)),
Err(e) => Err((e, usage)),
}
}
fn connection_to_sync_target(connection: ComposioConnection) -> Option<SyncTarget> {
@@ -321,6 +321,7 @@ pub(crate) async fn run_one_tick() -> Result<(), String> {
config: Arc::clone(&config),
toolkit: toolkit.clone(),
connection_id: Some(conn.id.clone()),
usage: Default::default(),
};
tracing::debug!(
@@ -284,7 +284,8 @@ pub use scope_lookup::{curated_scope_for, toolkit_has_scope};
pub use tool_scope::{classify_unknown, find_curated, toolkit_from_slug, CuratedTool, ToolScope};
pub use traits::ComposioProvider;
pub use types::{
NormalizedTask, ProviderContext, ProviderUserProfile, SyncOutcome, SyncReason, TaskFetchFilter,
ComposioUsage, ComposioUsageHandle, NormalizedTask, ProviderContext, ProviderUserProfile,
SyncOutcome, SyncReason, TaskFetchFilter,
};
pub use user_scopes::{load_or_default as load_user_scope_or_default, UserScopePref};
@@ -110,6 +110,7 @@ pub async fn sync_trigger_rpc(
config: Arc::clone(&config_arc),
toolkit: conn.toolkit.clone(),
connection_id: Some(conn.id.clone()),
usage: Default::default(),
};
match provider.sync(&ctx, SyncReason::Manual).await {
Ok(o) => outcomes.push(o),
@@ -1,7 +1,7 @@
//! Shared types for Composio provider implementations.
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use crate::openhuman::composio::client::{
create_composio_client, direct_execute, ComposioClient, ComposioClientKind,
@@ -196,11 +196,38 @@ impl TaskFetchFilter {
/// keeps an [`Arc<Config>`] and resolves the underlying client per call
/// through [`ProviderContext::execute`], mirroring the agent-tool
/// migration in [`crate::openhuman::composio::tools::ComposioExecuteTool`].
/// Per-sync accumulator for Composio billable-action usage.
///
/// Lives behind a shared handle on [`ProviderContext`] so the single
/// `execute` chokepoint can tally every action a provider fires during one
/// sync run, regardless of which provider (gmail / slack / github / notion /
/// linear / clickup) or how many pages it paginates.
/// [`crate::openhuman::memory_sync::composio::run_connection_sync`] returns
/// the final tally alongside the [`SyncOutcome`] for the sync audit log
/// (#3111).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ComposioUsage {
/// Count of `execute` calls that returned a response this run.
pub actions_called: u32,
/// Sum of each response's backend-reported `cost_usd`.
pub cost_usd: f64,
}
/// Shared, interior-mutable handle to a [`ComposioUsage`] tally. Cloning a
/// [`ProviderContext`] shares the same underlying counter, so the count is
/// stable no matter how the context is passed around within a sync.
pub type ComposioUsageHandle = Arc<Mutex<ComposioUsage>>;
#[derive(Clone)]
pub struct ProviderContext {
pub config: Arc<Config>,
pub toolkit: String,
pub connection_id: Option<String>,
/// Accumulates Composio billable-action usage across this context's
/// lifetime. Defaulted at every construction site; only the sync path
/// (`run_connection_sync`) reads it back. Non-sync callers (agent tools,
/// task-source fetches) leave it at zero — harmless.
pub usage: ComposioUsageHandle,
}
impl ProviderContext {
@@ -230,6 +257,7 @@ impl ProviderContext {
config,
toolkit: toolkit.into(),
connection_id,
usage: ComposioUsageHandle::default(),
}),
Err(e) => {
tracing::debug!(
@@ -282,7 +310,7 @@ impl ProviderContext {
anyhow::anyhow!("composio provider_context: failed to reload live config: {e}")
})?;
let kind = create_composio_client(&live_config)?;
match kind {
let result = match kind {
ComposioClientKind::Backend(client) => {
tracing::debug!(
action = %action,
@@ -299,7 +327,21 @@ impl ProviderContext {
);
direct_execute(&direct, action, arguments, &live_config.composio.entity_id).await
}
};
// Tally billable-action usage at the single chokepoint every provider
// routes through (#3111). We count any *completed* round-trip — even a
// provider-reported failure (`successful == false`) is a billable call
// — and sum the backend-reported `cost_usd`. Transport errors (the
// `Err` arm) never reached Composio, so they don't count. The lock is
// held only for the increment, never across an `.await`.
if let Ok(ref resp) = result {
if let Ok(mut usage) = self.usage.lock() {
usage.actions_called = usage.actions_called.saturating_add(1);
usage.cost_usd += resp.cost_usd;
}
}
result
}
/// Resolve a `ComposioClient` for callers that need a handle to
@@ -364,6 +406,46 @@ impl ProviderContext {
mod tests {
use super::*;
/// The whole #3111 tally relies on the `usage` handle being *shared*
/// across `ProviderContext` clones: a provider's `sync` runs against a
/// clone (or the same ctx passed by `&`), accumulates via `execute`, and
/// `run_connection_sync` reads the count back from its own handle. Pin
/// that the `Arc<Mutex<_>>` is genuinely shared so a clone's increments
/// are visible from the original — if this regressed to a per-clone
/// counter, the audit cost would silently always read zero.
#[test]
fn usage_handle_is_shared_across_context_clones() {
let ctx = ProviderContext {
config: Arc::new(Config::default()),
toolkit: "gmail".to_string(),
connection_id: None,
usage: ComposioUsageHandle::default(),
};
let cloned = ctx.clone();
// Simulate two `execute` round-trips accumulating on the clone.
{
let mut usage = cloned.usage.lock().expect("lock usage");
usage.actions_called = usage.actions_called.saturating_add(2);
usage.cost_usd += 0.015;
}
// The original handle must observe the clone's tally.
let observed = ctx.usage.lock().expect("lock usage");
assert_eq!(observed.actions_called, 2);
assert!((observed.cost_usd - 0.015).abs() < 1e-9);
}
/// `ComposioUsage` defaults to a zero tally — the value
/// `run_connection_sync` returns for a sync that fired no Composio
/// actions, and what non-sync `ProviderContext` callers carry.
#[test]
fn composio_usage_defaults_to_zero() {
let usage = ComposioUsage::default();
assert_eq!(usage.actions_called, 0);
assert_eq!(usage.cost_usd, 0.0);
}
// `ProviderContext::execute` and `ProviderContext::backend_client` reload
// config from `ctx.config.config_path` (via `reload_config_snapshot_with_timeout`)
// rather than from the process-global `OPENHUMAN_WORKSPACE`. Tests
@@ -392,6 +474,7 @@ mod tests {
config: Arc::new(config),
toolkit: "gmail".to_string(),
connection_id: None,
usage: ComposioUsageHandle::default(),
};
let res = ctx.execute("GMAIL_FETCH_EMAILS", None).await;
// The actual HTTP call will fail in the unit-test sandbox, but
@@ -424,6 +507,7 @@ mod tests {
config: Arc::new(config),
toolkit: "gmail".to_string(),
connection_id: None,
usage: ComposioUsageHandle::default(),
};
let res = ctx.execute("GMAIL_FETCH_EMAILS", None).await;
let err = res.expect_err("no backend session must error");
+58 -7
View File
@@ -34,6 +34,17 @@ pub struct SyncAuditEntry {
/// where the backend reported no charge — still render a cost. Prefer
/// [`SyncAuditEntry::actual_charged_usd`] when it is `Some`.
pub estimated_cost_usd: f64,
/// Number of Composio billable API actions executed during this sync
/// (e.g. `GMAIL_FETCH_EMAILS`, `SLACK_LIST_CONVERSATIONS`). `0` for
/// non-Composio source kinds. `#[serde(default)]` keeps audit lines
/// written before #3111 parseable.
#[serde(default)]
pub composio_actions_called: u32,
/// Actual USD charged for those Composio actions, summed from each
/// response's backend-reported `cost_usd`. `0.0` for non-Composio kinds
/// or when the backend reports no charge (e.g. direct mode).
#[serde(default)]
pub composio_cost_usd: f64,
/// Real amount billed by the backend in USD (sum of
/// `openhuman.billing.charged_amount_usd` across batches), when the
/// provider reported it for the run. `None` for runs that fell back to
@@ -52,6 +63,16 @@ pub struct SyncAuditEntry {
pub error: Option<String>,
}
impl SyncAuditEntry {
/// Total cost of the run: LLM summarisation cost plus the actual
/// Composio API-action cost. This is the "combined cost" the Sync
/// History UI surfaces so users see the full expense of a sync in one
/// number rather than just the summarisation slice (#3111).
pub fn combined_cost_usd(&self) -> f64 {
self.estimated_cost_usd + self.composio_cost_usd
}
}
const AUDIT_FILENAME: &str = "sync_audit.jsonl";
/// Append an audit entry to the sync audit log.
@@ -329,6 +350,8 @@ mod tests {
input_tokens: 50_000,
output_tokens: 5_000,
estimated_cost_usd: 0.225,
composio_actions_called: 0,
composio_cost_usd: 0.0,
actual_charged_usd: None,
duration_ms: 12_000,
success: true,
@@ -355,6 +378,8 @@ mod tests {
input_tokens: 100,
output_tokens: 10,
estimated_cost_usd: estimated,
composio_actions_called: 0,
composio_cost_usd: 0.0,
actual_charged_usd: actual,
duration_ms: 1,
success: true,
@@ -362,11 +387,41 @@ mod tests {
}
}
#[test]
fn combined_cost_sums_llm_and_composio() {
let entry = SyncAuditEntry {
timestamp: Utc::now(),
source_id: "src_gmail".to_string(),
source_kind: "composio".to_string(),
scope: "gmail".to_string(),
items_fetched: 40,
batches: 1,
input_tokens: 20_000,
output_tokens: 2_000,
estimated_cost_usd: 0.001_96,
composio_actions_called: 8,
composio_cost_usd: 0.04,
actual_charged_usd: None,
duration_ms: 5_000,
success: true,
error: None,
};
assert!((entry.combined_cost_usd() - 0.041_96).abs() < 1e-9);
}
#[test]
fn legacy_audit_line_without_composio_fields_deserializes() {
let legacy = r#"{"timestamp":"2026-05-01T00:00:00Z","source_id":"src_old","source_kind":"github_repo","scope":"github:org/repo","items_fetched":10,"batches":1,"input_tokens":1000,"output_tokens":100,"estimated_cost_usd":0.0001,"duration_ms":2000,"success":true}"#;
let entry: SyncAuditEntry =
serde_json::from_str(legacy).expect("legacy audit line must still parse");
assert_eq!(entry.composio_actions_called, 0);
assert_eq!(entry.composio_cost_usd, 0.0);
assert_eq!(entry.actual_charged_usd, None);
assert_eq!(entry.source_id, "src_old");
}
#[test]
fn effective_cost_prefers_actual_charge_when_present() {
// An entry built from a provider `UsageInfo` carries the real
// backend charge — `effective_cost_usd` must return it, not the
// hardcoded-pricing estimate.
let entry = entry_with_costs(0.0049, Some(0.0123));
assert!(entry.cost_is_actual());
assert!((entry.effective_cost_usd() - 0.0123).abs() < f64::EPSILON);
@@ -374,7 +429,6 @@ mod tests {
#[test]
fn effective_cost_falls_back_to_estimate_without_usage() {
// No provider usage → fall back to the estimate.
let entry = entry_with_costs(0.0049, None);
assert!(!entry.cost_is_actual());
assert!((entry.effective_cost_usd() - 0.0049).abs() < f64::EPSILON);
@@ -382,9 +436,6 @@ mod tests {
#[test]
fn old_entry_without_actual_field_deserializes_and_renders_estimate() {
// A pre-#3110 audit line has no `actual_charged_usd` key. The
// `#[serde(default)]` must let it deserialize, and the entry must
// render its estimate via `effective_cost_usd`.
let legacy = r#"{
"timestamp": "2024-01-01T00:00:00Z",
"source_id": "src_old",
@@ -322,6 +322,8 @@ pub async fn run_github_sync(
input_tokens: audit_input_tokens,
output_tokens: audit_output_tokens,
estimated_cost_usd: estimated_cost,
composio_actions_called: 0,
composio_cost_usd: 0.0,
actual_charged_usd,
duration_ms,
success: true,
@@ -274,6 +274,8 @@ pub async fn rebuild_tree_from_raw(config: &Config, scope: &str) -> Result<Rebui
input_tokens: audit_input_tokens,
output_tokens: audit_output_tokens,
estimated_cost_usd: estimated_cost,
composio_actions_called: 0,
composio_cost_usd: 0.0,
actual_charged_usd,
duration_ms,
success: true,
+1
View File
@@ -152,6 +152,7 @@ pub async fn preview_filter(
config: Arc::new(config.clone()),
toolkit: provider.as_str().to_string(),
connection_id: connection_id.filter(|s| !s.trim().is_empty()),
usage: Default::default(),
};
let max = max.unwrap_or(config.task_sources.max_tasks_per_fetch);
let fetch_filter = filter::to_fetch_filter(&filter_spec, max);
+1
View File
@@ -101,6 +101,7 @@ async fn run_inner(
config: Arc::new(config.clone()),
toolkit: source.provider.as_str().to_string(),
connection_id: source.connection_id.clone(),
usage: Default::default(),
};
let fetch_filter = filter::to_fetch_filter(&source.filter, source.max_tasks_per_fetch);
@@ -299,6 +299,7 @@ async fn configured_loopback_context(
config: Arc::new(config.clone()),
toolkit: toolkit.to_string(),
connection_id: Some(connection_id.to_string()),
usage: Default::default(),
};
(config, ctx, server)
}
@@ -550,6 +551,7 @@ async fn github_clickup_and_composio_bus_cover_provider_branches() {
config: github_ctx.config.clone(),
toolkit: "clickup".to_string(),
connection_id: Some("conn-clickup-round17".to_string()),
usage: Default::default(),
};
let clickup = ClickUpProvider::new();
let click_profile = clickup
@@ -164,6 +164,7 @@ async fn configured_context(
config: Arc::new(config.clone()),
toolkit: toolkit.to_string(),
connection_id: Some(connection_id.to_string()),
usage: Default::default(),
};
(config, ctx, server)
}
@@ -259,6 +259,7 @@ async fn configured_loopback_context(
config: Arc::new(config.clone()),
toolkit: "slack".to_string(),
connection_id: Some("conn-slack-round19".to_string()),
usage: Default::default(),
};
(config, ctx, server)
}
@@ -507,6 +507,7 @@ async fn composio_providers_fetch_profiles_tasks_and_cover_error_branches() {
config: Arc::new(config.clone()),
toolkit: "github".to_string(),
connection_id: Some("conn-github".to_string()),
usage: Default::default(),
};
let github = GitHubProvider::new();
let github_profile = github
@@ -238,6 +238,7 @@ async fn configured_loopback_context(
config: Arc::new(config.clone()),
toolkit: toolkit.to_string(),
connection_id: Some(connection_id.to_string()),
usage: Default::default(),
};
(config, ctx, server)
}
+1
View File
@@ -3043,6 +3043,7 @@ async fn memory_sync_provider_trait_defaults_and_connection_hook_are_determinist
config: Arc::new(config_in(&tmp)),
toolkit: "raw_coverage".into(),
connection_id: Some("conn-1".into()),
usage: Default::default(),
};
let provider = RawCoverageProvider { fail_profile: true };
assert_eq!(provider.sync_interval_secs(), Some(15 * 60));
@@ -598,6 +598,7 @@ async fn default_composio_provider_hooks_return_expected_noop_shapes() {
config: cfg,
toolkit: "round14".into(),
connection_id: Some("conn-round14".into()),
usage: Default::default(),
};
let provider = MinimalProvider;
assert_eq!(provider.sync_interval_secs(), None);