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::db;
12use crate::local::hooks::task_board_context::{TaskBoardContextHook, TaskBoardContextHookOptions};
13use crate::models::AgentState;
14use crate::stakpak::{StakpakApiClient, StakpakApiConfig};
15use libsql::Connection;
16use stakpak_shared::hooks::{HookRegistry, LifecycleEvent};
17use stakpak_shared::models::llm::{LLMModel, LLMProviderConfig, ProviderConfig};
18use stakpak_shared::models::stakai_adapter::{StakAIClient, get_stakai_model_string};
19use std::path::PathBuf;
20use std::sync::Arc;
21
22// =============================================================================
23// AgentClient Configuration
24// =============================================================================
25
26/// Model options for the AgentClient
27#[derive(Clone, Debug, Default)]
28pub struct ModelOptions {
29    /// Primary model for complex tasks
30    pub smart_model: Option<LLMModel>,
31    /// Economy model for simpler tasks
32    pub eco_model: Option<LLMModel>,
33    /// Fallback model when primary providers fail
34    pub recovery_model: Option<LLMModel>,
35}
36
37/// Default Stakpak API endpoint
38pub const DEFAULT_STAKPAK_ENDPOINT: &str = "https://apiv2.stakpak.dev";
39
40/// Stakpak connection configuration
41#[derive(Debug, Clone)]
42pub struct StakpakConfig {
43    /// Stakpak API key
44    pub api_key: String,
45    /// Stakpak API endpoint (default: https://apiv2.stakpak.dev)
46    pub api_endpoint: String,
47}
48
49impl StakpakConfig {
50    pub fn new(api_key: impl Into<String>) -> Self {
51        Self {
52            api_key: api_key.into(),
53            api_endpoint: DEFAULT_STAKPAK_ENDPOINT.to_string(),
54        }
55    }
56
57    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
58        self.api_endpoint = endpoint.into();
59        self
60    }
61}
62
63/// Configuration for creating an AgentClient
64#[derive(Debug, Default)]
65pub struct AgentClientConfig {
66    /// Stakpak configuration (optional - enables remote features when present)
67    pub stakpak: Option<StakpakConfig>,
68    /// LLM provider configurations
69    pub providers: LLMProviderConfig,
70    /// Smart model override
71    pub smart_model: Option<String>,
72    /// Eco model override
73    pub eco_model: Option<String>,
74    /// Recovery model override
75    pub recovery_model: Option<String>,
76    /// Local database path (default: ~/.stakpak/data/local.db)
77    pub store_path: Option<String>,
78    /// Hook registry for lifecycle events
79    pub hook_registry: Option<HookRegistry<AgentState>>,
80}
81
82impl AgentClientConfig {
83    /// Create new config
84    pub fn new() -> Self {
85        Self::default()
86    }
87
88    /// Set Stakpak configuration
89    ///
90    /// Use `StakpakConfig::new(api_key).with_endpoint(endpoint)` to configure.
91    pub fn with_stakpak(mut self, config: StakpakConfig) -> Self {
92        self.stakpak = Some(config);
93        self
94    }
95
96    /// Set providers
97    pub fn with_providers(mut self, providers: LLMProviderConfig) -> Self {
98        self.providers = providers;
99        self
100    }
101
102    /// Set smart model
103    pub fn with_smart_model(mut self, model: impl Into<String>) -> Self {
104        self.smart_model = Some(model.into());
105        self
106    }
107
108    /// Set eco model
109    pub fn with_eco_model(mut self, model: impl Into<String>) -> Self {
110        self.eco_model = Some(model.into());
111        self
112    }
113
114    /// Set recovery model
115    pub fn with_recovery_model(mut self, model: impl Into<String>) -> Self {
116        self.recovery_model = Some(model.into());
117        self
118    }
119
120    /// Set local database path
121    pub fn with_store_path(mut self, path: impl Into<String>) -> Self {
122        self.store_path = Some(path.into());
123        self
124    }
125
126    /// Set hook registry
127    pub fn with_hook_registry(mut self, registry: HookRegistry<AgentState>) -> Self {
128        self.hook_registry = Some(registry);
129        self
130    }
131}
132
133// =============================================================================
134// AgentClient
135// =============================================================================
136
137const DEFAULT_STORE_PATH: &str = ".stakpak/data/local.db";
138
139/// Unified agent client
140///
141/// Provides a single interface for:
142/// - LLM inference via stakai (with Stakpak or direct providers)
143/// - Session/checkpoint management (Stakpak API or local SQLite)
144/// - MCP tools, billing, rulebooks (Stakpak API only)
145#[derive(Clone)]
146pub struct AgentClient {
147    /// StakAI client for all LLM inference
148    pub(crate) stakai: StakAIClient,
149    /// Stakpak API client for non-inference operations (optional)
150    pub(crate) stakpak_api: Option<StakpakApiClient>,
151    /// Local SQLite database for fallback storage
152    pub(crate) local_db: Connection,
153    /// Hook registry for lifecycle events
154    pub(crate) hook_registry: Arc<HookRegistry<AgentState>>,
155    /// Model configuration
156    pub(crate) model_options: ModelOptions,
157    /// Stakpak configuration (for reference)
158    pub(crate) stakpak: Option<StakpakConfig>,
159}
160
161impl AgentClient {
162    /// Create a new AgentClient
163    pub async fn new(config: AgentClientConfig) -> Result<Self, String> {
164        // 1. Build LLMProviderConfig with Stakpak if configured (only if api_key is not empty)
165        let mut providers = config.providers.clone();
166        if let Some(stakpak) = &config.stakpak
167            && !stakpak.api_key.is_empty()
168        {
169            providers.providers.insert(
170                "stakpak".to_string(),
171                ProviderConfig::Stakpak {
172                    api_key: stakpak.api_key.clone(),
173                    api_endpoint: Some(stakpak.api_endpoint.clone()),
174                },
175            );
176        }
177
178        // 2. Create StakAIClient with all providers
179        let stakai = StakAIClient::new(&providers)
180            .map_err(|e| format!("Failed to create StakAI client: {}", e))?;
181
182        // 3. Create StakpakApiClient if configured (only if api_key is not empty)
183        let stakpak_api = if let Some(stakpak) = &config.stakpak {
184            if !stakpak.api_key.is_empty() {
185                Some(
186                    StakpakApiClient::new(&StakpakApiConfig {
187                        api_key: stakpak.api_key.clone(),
188                        api_endpoint: stakpak.api_endpoint.clone(),
189                    })
190                    .map_err(|e| format!("Failed to create Stakpak API client: {}", e))?,
191                )
192            } else {
193                None
194            }
195        } else {
196            None
197        };
198
199        // 4. Initialize local SQLite database
200        let store_path = config.store_path.map(PathBuf::from).unwrap_or_else(|| {
201            std::env::var("HOME")
202                .map(PathBuf::from)
203                .unwrap_or_default()
204                .join(DEFAULT_STORE_PATH)
205        });
206
207        if let Some(parent) = store_path.parent() {
208            std::fs::create_dir_all(parent)
209                .map_err(|e| format!("Failed to create database directory: {}", e))?;
210        }
211
212        let db = libsql::Builder::new_local(store_path.display().to_string())
213            .build()
214            .await
215            .map_err(|e| format!("Failed to open database: {}", e))?;
216        let local_db = db
217            .connect()
218            .map_err(|e| format!("Failed to connect to database: {}", e))?;
219        db::init_schema(&local_db).await?;
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            local_db,
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 local database
279    pub fn local_db(&self) -> &Connection {
280        &self.local_db
281    }
282
283    /// Get reference to the hook registry
284    pub fn hook_registry(&self) -> &Arc<HookRegistry<AgentState>> {
285        &self.hook_registry
286    }
287
288    /// Get the model options
289    pub fn model_options(&self) -> &ModelOptions {
290        &self.model_options
291    }
292
293    /// Get the model string for the given agent model type
294    ///
295    /// When Stakpak is available, routes through Stakpak provider.
296    /// Otherwise, uses direct provider.
297    pub fn get_model_string(
298        &self,
299        model: &stakpak_shared::models::integrations::openai::AgentModel,
300    ) -> LLMModel {
301        use stakpak_shared::models::integrations::openai::AgentModel;
302
303        let base_model = match model {
304            AgentModel::Smart => self.model_options.smart_model.clone().unwrap_or_else(|| {
305                LLMModel::from("anthropic/claude-sonnet-4-5-20250929".to_string())
306            }),
307            AgentModel::Eco => self.model_options.eco_model.clone().unwrap_or_else(|| {
308                LLMModel::from("anthropic/claude-haiku-4-5-20250929".to_string())
309            }),
310            AgentModel::Recovery => self
311                .model_options
312                .recovery_model
313                .clone()
314                .unwrap_or_else(|| LLMModel::from("openai/gpt-5".to_string())),
315        };
316
317        // If Stakpak is available, route through Stakpak provider
318        if self.has_stakpak() {
319            // Get properly formatted model string with provider prefix (e.g., "anthropic/claude-sonnet-4-5")
320            let model_str = get_stakai_model_string(&base_model);
321            // Extract display name from the last segment for UI
322            let display_name = model_str
323                .rsplit('/')
324                .next()
325                .unwrap_or(&model_str)
326                .to_string();
327            LLMModel::Custom {
328                provider: "stakpak".to_string(),
329                model: model_str,
330                name: Some(display_name),
331            }
332        } else {
333            base_model
334        }
335    }
336}
337
338// Debug implementation for AgentClient
339impl std::fmt::Debug for AgentClient {
340    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
341        f.debug_struct("AgentClient")
342            .field("has_stakpak", &self.has_stakpak())
343            .field("model_options", &self.model_options)
344            .finish_non_exhaustive()
345    }
346}
stakpak_api/client/mod.rs

stakpak_api/client/
mod.rs
