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}