runtara_agents/connections.rs
1// Copyright (C) 2025 SyncMyOrders Sp. z o.o.
2// SPDX-License-Identifier: AGPL-3.0-or-later
3//! Shared connection management for agents
4//!
5//! This module provides agent-neutral connection storage and retrieval.
6//!
7//! Connections can come from two sources:
8//! 1. Injected via `_connection` field in input (primary - from connection service)
9//! 2. Thread-local storage (legacy - for testing service)
10//!
11//! Agents should use `resolve_connection()` which tries both sources.
12
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use std::cell::RefCell;
16use std::collections::HashMap;
17
18/// Raw connection data returned from host
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct RawConnection {
21 #[serde(skip_serializing_if = "Option::is_none")]
22 pub connection_subtype: Option<String>,
23 /// Connection type identifier that maps to a connection schema (e.g., bearer, api_key, sftp)
24 pub integration_id: String,
25 pub parameters: Value,
26 #[serde(skip_serializing_if = "Option::is_none")]
27 pub rate_limit_config: Option<Value>,
28}
29
30thread_local! {
31 /// Thread-local storage for connection data (legacy - used by testing service)
32 static CONNECTIONS: RefCell<HashMap<String, RawConnection>> = RefCell::new(HashMap::new());
33}
34
35/// Register a connection for the current thread (used by testing service)
36pub fn register_connection(connection_id: &str, connection: RawConnection) {
37 CONNECTIONS.with(|c| {
38 c.borrow_mut().insert(connection_id.to_string(), connection);
39 });
40}
41
42/// Clear all registered connections for the current thread
43pub fn clear_connections() {
44 CONNECTIONS.with(|c| {
45 c.borrow_mut().clear();
46 });
47}
48
49/// Get connection data by ID from thread-local storage (legacy)
50pub fn get_connection(connection_id: &str) -> Result<RawConnection, String> {
51 let result = CONNECTIONS.with(|c| c.borrow().get(connection_id).cloned());
52
53 match result {
54 Some(conn) => Ok(conn),
55 None => Err(format!(
56 "Connection '{}' not found in registered connections.",
57 connection_id
58 )),
59 }
60}
61
62/// Parse connection data from input JSON (provided by caller, e.g., agent testing service)
63pub fn parse_connection_from_input(input: &Value) -> Option<RawConnection> {
64 input
65 .get("_connection")
66 .and_then(|conn| serde_json::from_value(conn.clone()).ok())
67}
68
69/// Resolve a connection from the current capability input or thread-local storage.
70///
71/// This is the primary function agents should use to get connection data.
72/// It tries the following sources in order:
73/// 1. `_connection` field in the current capability input (set by executor wrapper)
74/// 2. Thread-local connections storage (legacy - for testing service)
75///
76/// # Arguments
77/// * `connection_id` - The connection identifier
78///
79/// # Returns
80/// The resolved connection or an error if not found in any source
81pub fn resolve_connection(connection_id: &str) -> Result<RawConnection, String> {
82 // Primary: check current capability input (set by executor wrapper)
83 if let Some(input) = runtara_dsl::agent_meta::get_current_input() {
84 if let Some(conn) = parse_connection_from_input(&input) {
85 return Ok(conn);
86 }
87 }
88
89 // Fallback: thread-local storage (legacy - testing service)
90 get_connection(connection_id)
91}