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::{LLMModel, LLMProviderConfig, ProviderConfig};
20use stakpak_shared::models::stakai_adapter::{StakAIClient, get_stakai_model_string};
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<LLMModel>,
32    /// Economy model for simpler tasks
33    pub eco_model: Option<LLMModel>,
34    /// Fallback model when primary providers fail
35    pub recovery_model: Option<LLMModel>,
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.map(LLMModel::from),
224            eco_model: config.eco_model.map(LLMModel::from),
225            recovery_model: config.recovery_model.map(LLMModel::from),
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                model_options: crate::local::ModelOptions {
234                    smart_model: model_options.smart_model.clone(),
235                    eco_model: model_options.eco_model.clone(),
236                    recovery_model: model_options.recovery_model.clone(),
237                },
238                history_action_message_size_limit: Some(100),
239                history_action_message_keep_last_n: Some(50),
240                history_action_result_keep_last_n: Some(50),
241            })),
242        );
243        let hook_registry = Arc::new(hook_registry);
244
245        Ok(Self {
246            stakai,
247            stakpak_api,
248            session_storage,
249            hook_registry,
250            model_options,
251            stakpak: config.stakpak,
252        })
253    }
254
255    /// Check if Stakpak API is available
256    pub fn has_stakpak(&self) -> bool {
257        self.stakpak_api.is_some()
258    }
259
260    /// Get the Stakpak API endpoint (with default fallback)
261    pub fn get_stakpak_api_endpoint(&self) -> &str {
262        self.stakpak
263            .as_ref()
264            .map(|s| s.api_endpoint.as_str())
265            .unwrap_or(DEFAULT_STAKPAK_ENDPOINT)
266    }
267
268    /// Get reference to the StakAI client
269    pub fn stakai(&self) -> &StakAIClient {
270        &self.stakai
271    }
272
273    /// Get reference to the Stakpak API client (if available)
274    pub fn stakpak_api(&self) -> Option<&StakpakApiClient> {
275        self.stakpak_api.as_ref()
276    }
277
278    /// Get reference to the hook registry
279    pub fn hook_registry(&self) -> &Arc<HookRegistry<AgentState>> {
280        &self.hook_registry
281    }
282
283    /// Get the model options
284    pub fn model_options(&self) -> &ModelOptions {
285        &self.model_options
286    }
287
288    /// Get reference to the session storage
289    ///
290    /// Use this for all session and checkpoint operations.
291    pub fn session_storage(&self) -> &Arc<dyn SessionStorage> {
292        &self.session_storage
293    }
294
295    /// Get the model string for the given agent model type
296    ///
297    /// When Stakpak is available, routes through Stakpak provider.
298    /// Otherwise, uses direct provider.
299    pub fn get_model_string(
300        &self,
301        model: &stakpak_shared::models::integrations::openai::AgentModel,
302    ) -> LLMModel {
303        use stakpak_shared::models::integrations::openai::AgentModel;
304
305        let base_model = match model {
306            AgentModel::Smart => self.model_options.smart_model.clone().unwrap_or_else(|| {
307                LLMModel::from("anthropic/claude-sonnet-4-5-20250929".to_string())
308            }),
309            AgentModel::Eco => self.model_options.eco_model.clone().unwrap_or_else(|| {
310                LLMModel::from("anthropic/claude-haiku-4-5-20250929".to_string())
311            }),
312            AgentModel::Recovery => self
313                .model_options
314                .recovery_model
315                .clone()
316                .unwrap_or_else(|| LLMModel::from("openai/gpt-5".to_string())),
317        };
318
319        // If Stakpak is available, route through Stakpak provider
320        if self.has_stakpak() {
321            // Get properly formatted model string with provider prefix (e.g., "anthropic/claude-sonnet-4-5")
322            let model_str = get_stakai_model_string(&base_model);
323            // Extract display name from the last segment for UI
324            let display_name = model_str
325                .rsplit('/')
326                .next()
327                .unwrap_or(&model_str)
328                .to_string();
329            LLMModel::Custom {
330                provider: "stakpak".to_string(),
331                model: model_str,
332                name: Some(display_name),
333            }
334        } else {
335            base_model
336        }
337    }
338}
339
340// Debug implementation for AgentClient
341impl std::fmt::Debug for AgentClient {
342    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
343        f.debug_struct("AgentClient")
344            .field("has_stakpak", &self.has_stakpak())
345            .field("model_options", &self.model_options)
346            .finish_non_exhaustive()
347    }
348}
stakpak_api/client/mod.rs

stakpak_api/client/
mod.rs
