stakpak_api/client/
mod.rs

1//! Unified AgentClient
2//!
3//! The AgentClient provides a unified interface that:
4//! - Uses stakai for all LLM inference (with StakpakProvider when available)
5//! - Uses StakpakApiClient for non-inference APIs (sessions, billing, etc.)
6//! - Falls back to local SQLite DB when Stakpak is unavailable
7//! - Integrates with hooks for lifecycle events
8
9mod provider;
10
11use crate::local::hooks::task_board_context::{TaskBoardContextHook, TaskBoardContextHookOptions};
12use crate::local::storage::LocalStorage;
13use crate::models::AgentState;
14use crate::stakpak::storage::StakpakStorage;
15use crate::stakpak::{StakpakApiClient, StakpakApiConfig};
16use crate::storage::SessionStorage;
17
18use stakpak_shared::hooks::{HookRegistry, LifecycleEvent};
19use stakpak_shared::models::llm::{LLMProviderConfig, ProviderConfig};
20use stakpak_shared::models::stakai_adapter::StakAIClient;
21use std::sync::Arc;
22
23// =============================================================================
24// AgentClient Configuration
25// =============================================================================
26
27/// Default Stakpak API endpoint
28pub const DEFAULT_STAKPAK_ENDPOINT: &str = "https://apiv2.stakpak.dev";
29
30/// Stakpak connection configuration
31#[derive(Debug, Clone)]
32pub struct StakpakConfig {
33    /// Stakpak API key
34    pub api_key: String,
35    /// Stakpak API endpoint (default: https://apiv2.stakpak.dev)
36    pub api_endpoint: String,
37}
38
39impl StakpakConfig {
40    pub fn new(api_key: impl Into<String>) -> Self {
41        Self {
42            api_key: api_key.into(),
43            api_endpoint: DEFAULT_STAKPAK_ENDPOINT.to_string(),
44        }
45    }
46
47    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
48        self.api_endpoint = endpoint.into();
49        self
50    }
51}
52
53/// Configuration for creating an AgentClient
54#[derive(Debug, Default)]
55pub struct AgentClientConfig {
56    /// Stakpak configuration (optional - enables remote features when present)
57    pub stakpak: Option<StakpakConfig>,
58    /// LLM provider configurations
59    pub providers: LLMProviderConfig,
60    /// Local database path (default: ~/.stakpak/data/local.db)
61    pub store_path: Option<String>,
62    /// Hook registry for lifecycle events
63    pub hook_registry: Option<HookRegistry<AgentState>>,
64}
65
66impl AgentClientConfig {
67    /// Create new config
68    pub fn new() -> Self {
69        Self::default()
70    }
71
72    /// Set Stakpak configuration
73    ///
74    /// Use `StakpakConfig::new(api_key).with_endpoint(endpoint)` to configure.
75    pub fn with_stakpak(mut self, config: StakpakConfig) -> Self {
76        self.stakpak = Some(config);
77        self
78    }
79
80    /// Set providers
81    pub fn with_providers(mut self, providers: LLMProviderConfig) -> Self {
82        self.providers = providers;
83        self
84    }
85
86    /// Set local database path
87    pub fn with_store_path(mut self, path: impl Into<String>) -> Self {
88        self.store_path = Some(path.into());
89        self
90    }
91
92    /// Set hook registry
93    pub fn with_hook_registry(mut self, registry: HookRegistry<AgentState>) -> Self {
94        self.hook_registry = Some(registry);
95        self
96    }
97}
98
99// =============================================================================
100// AgentClient
101// =============================================================================
102
103const DEFAULT_STORE_PATH: &str = ".stakpak/data/local.db";
104
105/// Unified agent client
106///
107/// Provides a single interface for:
108/// - LLM inference via stakai (with Stakpak or direct providers)
109/// - Session/checkpoint management via SessionStorage trait (Stakpak API or local SQLite)
110/// - MCP tools, billing, rulebooks (Stakpak API only)
111#[derive(Clone)]
112pub struct AgentClient {
113    /// StakAI client for all LLM inference
114    pub(crate) stakai: StakAIClient,
115    /// Stakpak API client for non-inference operations (optional)
116    pub(crate) stakpak_api: Option<StakpakApiClient>,
117    /// Session storage implementation (abstracts Stakpak API vs local SQLite)
118    pub(crate) session_storage: Arc<dyn SessionStorage>,
119    /// Hook registry for lifecycle events
120    pub(crate) hook_registry: Arc<HookRegistry<AgentState>>,
121    /// Stakpak configuration (for reference)
122    pub(crate) stakpak: Option<StakpakConfig>,
123}
124
125impl AgentClient {
126    /// Create a new AgentClient
127    pub async fn new(config: AgentClientConfig) -> Result<Self, String> {
128        // 1. Build LLMProviderConfig with Stakpak if configured (only if api_key is not empty)
129        let mut providers = config.providers.clone();
130        if let Some(stakpak) = &config.stakpak
131            && !stakpak.api_key.is_empty()
132        {
133            providers.providers.insert(
134                "stakpak".to_string(),
135                ProviderConfig::Stakpak {
136                    api_key: Some(stakpak.api_key.clone()),
137                    api_endpoint: Some(stakpak.api_endpoint.clone()),
138                    auth: None,
139                },
140            );
141        }
142
143        // 2. Create StakAIClient with all providers
144        let stakai = StakAIClient::new(&providers)
145            .map_err(|e| format!("Failed to create StakAI client: {}", e))?;
146
147        // 3. Create StakpakApiClient if configured (only if api_key is not empty)
148        let stakpak_api = if let Some(stakpak) = &config.stakpak {
149            if !stakpak.api_key.is_empty() {
150                Some(
151                    StakpakApiClient::new(&StakpakApiConfig {
152                        api_key: stakpak.api_key.clone(),
153                        api_endpoint: stakpak.api_endpoint.clone(),
154                    })
155                    .map_err(|e| format!("Failed to create Stakpak API client: {}", e))?,
156                )
157            } else {
158                None
159            }
160        } else {
161            None
162        };
163
164        // 4. Create session storage (Stakpak API or local SQLite)
165        let session_storage: Arc<dyn SessionStorage> = if let Some(stakpak) = &config.stakpak
166            && !stakpak.api_key.is_empty()
167        {
168            Arc::new(
169                StakpakStorage::new(&stakpak.api_key, &stakpak.api_endpoint)
170                    .map_err(|e| format!("Failed to create Stakpak storage: {}", e))?,
171            )
172        } else {
173            let store_path = config.store_path.clone().unwrap_or_else(|| {
174                std::env::var("HOME")
175                    .map(|h| format!("{}/{}", h, DEFAULT_STORE_PATH))
176                    .unwrap_or_else(|_| DEFAULT_STORE_PATH.to_string())
177            });
178            Arc::new(
179                LocalStorage::new(&store_path)
180                    .await
181                    .map_err(|e| format!("Failed to create local storage: {}", e))?,
182            )
183        };
184
185        // 6. Setup hook registry with context management hooks
186        let mut hook_registry = config.hook_registry.unwrap_or_default();
187        hook_registry.register(
188            LifecycleEvent::BeforeInference,
189            Box::new(TaskBoardContextHook::new(TaskBoardContextHookOptions {
190                keep_last_n_assistant_messages: Some(5), // Keep the last 5 assistant messages in context
191                context_budget_threshold: Some(0.8),     // defaults to 0.8 (80%)
192            })),
193        );
194        let hook_registry = Arc::new(hook_registry);
195
196        Ok(Self {
197            stakai,
198            stakpak_api,
199            session_storage,
200            hook_registry,
201            stakpak: config.stakpak,
202        })
203    }
204
205    /// Check if Stakpak API is available
206    pub fn has_stakpak(&self) -> bool {
207        self.stakpak_api.is_some()
208    }
209
210    /// Get the Stakpak API endpoint (with default fallback)
211    pub fn get_stakpak_api_endpoint(&self) -> &str {
212        self.stakpak
213            .as_ref()
214            .map(|s| s.api_endpoint.as_str())
215            .unwrap_or(DEFAULT_STAKPAK_ENDPOINT)
216    }
217
218    /// Get reference to the StakAI client
219    pub fn stakai(&self) -> &StakAIClient {
220        &self.stakai
221    }
222
223    /// Get reference to the Stakpak API client (if available)
224    pub fn stakpak_api(&self) -> Option<&StakpakApiClient> {
225        self.stakpak_api.as_ref()
226    }
227
228    /// Get reference to the hook registry
229    pub fn hook_registry(&self) -> &Arc<HookRegistry<AgentState>> {
230        &self.hook_registry
231    }
232
233    /// Get reference to the session storage
234    ///
235    /// Use this for all session and checkpoint operations.
236    pub fn session_storage(&self) -> &Arc<dyn SessionStorage> {
237        &self.session_storage
238    }
239}
240
241// Debug implementation for AgentClient
242impl std::fmt::Debug for AgentClient {
243    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
244        f.debug_struct("AgentClient")
245            .field("has_stakpak", &self.has_stakpak())
246            .finish_non_exhaustive()
247    }
248}
stakpak_api/client/mod.rs

stakpak_api/client/
mod.rs
