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/// Model options for the AgentClient
28#[derive(Clone, Debug, Default)]
29pub struct ModelOptions {
30    /// Primary model for complex tasks
31    pub smart_model: Option<String>,
32    /// Economy model for simpler tasks
33    pub eco_model: Option<String>,
34    /// Fallback model when primary providers fail
35    pub recovery_model: Option<String>,
36}
37
38/// Default Stakpak API endpoint
39pub const DEFAULT_STAKPAK_ENDPOINT: &str = "https://apiv2.stakpak.dev";
40
41/// Stakpak connection configuration
42#[derive(Debug, Clone)]
43pub struct StakpakConfig {
44    /// Stakpak API key
45    pub api_key: String,
46    /// Stakpak API endpoint (default: https://apiv2.stakpak.dev)
47    pub api_endpoint: String,
48}
49
50impl StakpakConfig {
51    pub fn new(api_key: impl Into<String>) -> Self {
52        Self {
53            api_key: api_key.into(),
54            api_endpoint: DEFAULT_STAKPAK_ENDPOINT.to_string(),
55        }
56    }
57
58    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
59        self.api_endpoint = endpoint.into();
60        self
61    }
62}
63
64/// Configuration for creating an AgentClient
65#[derive(Debug, Default)]
66pub struct AgentClientConfig {
67    /// Stakpak configuration (optional - enables remote features when present)
68    pub stakpak: Option<StakpakConfig>,
69    /// LLM provider configurations
70    pub providers: LLMProviderConfig,
71    /// Smart model override
72    pub smart_model: Option<String>,
73    /// Eco model override
74    pub eco_model: Option<String>,
75    /// Recovery model override
76    pub recovery_model: Option<String>,
77    /// Local database path (default: ~/.stakpak/data/local.db)
78    pub store_path: Option<String>,
79    /// Hook registry for lifecycle events
80    pub hook_registry: Option<HookRegistry<AgentState>>,
81}
82
83impl AgentClientConfig {
84    /// Create new config
85    pub fn new() -> Self {
86        Self::default()
87    }
88
89    /// Set Stakpak configuration
90    ///
91    /// Use `StakpakConfig::new(api_key).with_endpoint(endpoint)` to configure.
92    pub fn with_stakpak(mut self, config: StakpakConfig) -> Self {
93        self.stakpak = Some(config);
94        self
95    }
96
97    /// Set providers
98    pub fn with_providers(mut self, providers: LLMProviderConfig) -> Self {
99        self.providers = providers;
100        self
101    }
102
103    /// Set smart model
104    pub fn with_smart_model(mut self, model: impl Into<String>) -> Self {
105        self.smart_model = Some(model.into());
106        self
107    }
108
109    /// Set eco model
110    pub fn with_eco_model(mut self, model: impl Into<String>) -> Self {
111        self.eco_model = Some(model.into());
112        self
113    }
114
115    /// Set recovery model
116    pub fn with_recovery_model(mut self, model: impl Into<String>) -> Self {
117        self.recovery_model = Some(model.into());
118        self
119    }
120
121    /// Set local database path
122    pub fn with_store_path(mut self, path: impl Into<String>) -> Self {
123        self.store_path = Some(path.into());
124        self
125    }
126
127    /// Set hook registry
128    pub fn with_hook_registry(mut self, registry: HookRegistry<AgentState>) -> Self {
129        self.hook_registry = Some(registry);
130        self
131    }
132}
133
134// =============================================================================
135// AgentClient
136// =============================================================================
137
138const DEFAULT_STORE_PATH: &str = ".stakpak/data/local.db";
139
140/// Unified agent client
141///
142/// Provides a single interface for:
143/// - LLM inference via stakai (with Stakpak or direct providers)
144/// - Session/checkpoint management via SessionStorage trait (Stakpak API or local SQLite)
145/// - MCP tools, billing, rulebooks (Stakpak API only)
146#[derive(Clone)]
147pub struct AgentClient {
148    /// StakAI client for all LLM inference
149    pub(crate) stakai: StakAIClient,
150    /// Stakpak API client for non-inference operations (optional)
151    pub(crate) stakpak_api: Option<StakpakApiClient>,
152    /// Session storage implementation (abstracts Stakpak API vs local SQLite)
153    pub(crate) session_storage: Arc<dyn SessionStorage>,
154    /// Hook registry for lifecycle events
155    pub(crate) hook_registry: Arc<HookRegistry<AgentState>>,
156    /// Model configuration
157    pub(crate) model_options: ModelOptions,
158    /// Stakpak configuration (for reference)
159    pub(crate) stakpak: Option<StakpakConfig>,
160}
161
162impl AgentClient {
163    /// Create a new AgentClient
164    pub async fn new(config: AgentClientConfig) -> Result<Self, String> {
165        // 1. Build LLMProviderConfig with Stakpak if configured (only if api_key is not empty)
166        let mut providers = config.providers.clone();
167        if let Some(stakpak) = &config.stakpak
168            && !stakpak.api_key.is_empty()
169        {
170            providers.providers.insert(
171                "stakpak".to_string(),
172                ProviderConfig::Stakpak {
173                    api_key: stakpak.api_key.clone(),
174                    api_endpoint: Some(stakpak.api_endpoint.clone()),
175                },
176            );
177        }
178
179        // 2. Create StakAIClient with all providers
180        let stakai = StakAIClient::new(&providers)
181            .map_err(|e| format!("Failed to create StakAI client: {}", e))?;
182
183        // 3. Create StakpakApiClient if configured (only if api_key is not empty)
184        let stakpak_api = if let Some(stakpak) = &config.stakpak {
185            if !stakpak.api_key.is_empty() {
186                Some(
187                    StakpakApiClient::new(&StakpakApiConfig {
188                        api_key: stakpak.api_key.clone(),
189                        api_endpoint: stakpak.api_endpoint.clone(),
190                    })
191                    .map_err(|e| format!("Failed to create Stakpak API client: {}", e))?,
192                )
193            } else {
194                None
195            }
196        } else {
197            None
198        };
199
200        // 4. Create session storage (Stakpak API or local SQLite)
201        let session_storage: Arc<dyn SessionStorage> = if let Some(stakpak) = &config.stakpak
202            && !stakpak.api_key.is_empty()
203        {
204            Arc::new(
205                StakpakStorage::new(&stakpak.api_key, &stakpak.api_endpoint)
206                    .map_err(|e| format!("Failed to create Stakpak storage: {}", e))?,
207            )
208        } else {
209            let store_path = config.store_path.clone().unwrap_or_else(|| {
210                std::env::var("HOME")
211                    .map(|h| format!("{}/{}", h, DEFAULT_STORE_PATH))
212                    .unwrap_or_else(|_| DEFAULT_STORE_PATH.to_string())
213            });
214            Arc::new(
215                LocalStorage::new(&store_path)
216                    .await
217                    .map_err(|e| format!("Failed to create local storage: {}", e))?,
218            )
219        };
220
221        // 5. Parse model options
222        let model_options = ModelOptions {
223            smart_model: config.smart_model,
224            eco_model: config.eco_model,
225            recovery_model: config.recovery_model,
226        };
227
228        // 6. Setup hook registry with context management hooks
229        let mut hook_registry = config.hook_registry.unwrap_or_default();
230        hook_registry.register(
231            LifecycleEvent::BeforeInference,
232            Box::new(TaskBoardContextHook::new(TaskBoardContextHookOptions {
233                keep_last_n_assistant_messages: Some(5), // Keep the last 5 assistant messages in context
234                context_budget_threshold: Some(0.8),     // defaults to 0.8 (80%)
235            })),
236        );
237        let hook_registry = Arc::new(hook_registry);
238
239        Ok(Self {
240            stakai,
241            stakpak_api,
242            session_storage,
243            hook_registry,
244            model_options,
245            stakpak: config.stakpak,
246        })
247    }
248
249    /// Check if Stakpak API is available
250    pub fn has_stakpak(&self) -> bool {
251        self.stakpak_api.is_some()
252    }
253
254    /// Get the Stakpak API endpoint (with default fallback)
255    pub fn get_stakpak_api_endpoint(&self) -> &str {
256        self.stakpak
257            .as_ref()
258            .map(|s| s.api_endpoint.as_str())
259            .unwrap_or(DEFAULT_STAKPAK_ENDPOINT)
260    }
261
262    /// Get reference to the StakAI client
263    pub fn stakai(&self) -> &StakAIClient {
264        &self.stakai
265    }
266
267    /// Get reference to the Stakpak API client (if available)
268    pub fn stakpak_api(&self) -> Option<&StakpakApiClient> {
269        self.stakpak_api.as_ref()
270    }
271
272    /// Get reference to the hook registry
273    pub fn hook_registry(&self) -> &Arc<HookRegistry<AgentState>> {
274        &self.hook_registry
275    }
276
277    /// Get the model options
278    pub fn model_options(&self) -> &ModelOptions {
279        &self.model_options
280    }
281
282    /// Get reference to the session storage
283    ///
284    /// Use this for all session and checkpoint operations.
285    pub fn session_storage(&self) -> &Arc<dyn SessionStorage> {
286        &self.session_storage
287    }
288}
289
290// Debug implementation for AgentClient
291impl std::fmt::Debug for AgentClient {
292    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293        f.debug_struct("AgentClient")
294            .field("has_stakpak", &self.has_stakpak())
295            .field("model_options", &self.model_options)
296            .finish_non_exhaustive()
297    }
298}
stakpak_api/client/mod.rs

stakpak_api/client/
mod.rs
