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: Some(stakpak.api_key.clone()),
174                    api_endpoint: Some(stakpak.api_endpoint.clone()),
175                    auth: None,
176                },
177            );
178        }
179
180        // 2. Create StakAIClient with all providers
181        let stakai = StakAIClient::new(&providers)
182            .map_err(|e| format!("Failed to create StakAI client: {}", e))?;
183
184        // 3. Create StakpakApiClient if configured (only if api_key is not empty)
185        let stakpak_api = if let Some(stakpak) = &config.stakpak {
186            if !stakpak.api_key.is_empty() {
187                Some(
188                    StakpakApiClient::new(&StakpakApiConfig {
189                        api_key: stakpak.api_key.clone(),
190                        api_endpoint: stakpak.api_endpoint.clone(),
191                    })
192                    .map_err(|e| format!("Failed to create Stakpak API client: {}", e))?,
193                )
194            } else {
195                None
196            }
197        } else {
198            None
199        };
200
201        // 4. Create session storage (Stakpak API or local SQLite)
202        let session_storage: Arc<dyn SessionStorage> = if let Some(stakpak) = &config.stakpak
203            && !stakpak.api_key.is_empty()
204        {
205            Arc::new(
206                StakpakStorage::new(&stakpak.api_key, &stakpak.api_endpoint)
207                    .map_err(|e| format!("Failed to create Stakpak storage: {}", e))?,
208            )
209        } else {
210            let store_path = config.store_path.clone().unwrap_or_else(|| {
211                std::env::var("HOME")
212                    .map(|h| format!("{}/{}", h, DEFAULT_STORE_PATH))
213                    .unwrap_or_else(|_| DEFAULT_STORE_PATH.to_string())
214            });
215            Arc::new(
216                LocalStorage::new(&store_path)
217                    .await
218                    .map_err(|e| format!("Failed to create local storage: {}", e))?,
219            )
220        };
221
222        // 5. Parse model options
223        let model_options = ModelOptions {
224            smart_model: config.smart_model,
225            eco_model: config.eco_model,
226            recovery_model: config.recovery_model,
227        };
228
229        // 6. Setup hook registry with context management hooks
230        let mut hook_registry = config.hook_registry.unwrap_or_default();
231        hook_registry.register(
232            LifecycleEvent::BeforeInference,
233            Box::new(TaskBoardContextHook::new(TaskBoardContextHookOptions {
234                keep_last_n_assistant_messages: Some(5), // Keep the last 5 assistant messages in context
235                context_budget_threshold: Some(0.8),     // defaults to 0.8 (80%)
236            })),
237        );
238        let hook_registry = Arc::new(hook_registry);
239
240        Ok(Self {
241            stakai,
242            stakpak_api,
243            session_storage,
244            hook_registry,
245            model_options,
246            stakpak: config.stakpak,
247        })
248    }
249
250    /// Check if Stakpak API is available
251    pub fn has_stakpak(&self) -> bool {
252        self.stakpak_api.is_some()
253    }
254
255    /// Get the Stakpak API endpoint (with default fallback)
256    pub fn get_stakpak_api_endpoint(&self) -> &str {
257        self.stakpak
258            .as_ref()
259            .map(|s| s.api_endpoint.as_str())
260            .unwrap_or(DEFAULT_STAKPAK_ENDPOINT)
261    }
262
263    /// Get reference to the StakAI client
264    pub fn stakai(&self) -> &StakAIClient {
265        &self.stakai
266    }
267
268    /// Get reference to the Stakpak API client (if available)
269    pub fn stakpak_api(&self) -> Option<&StakpakApiClient> {
270        self.stakpak_api.as_ref()
271    }
272
273    /// Get reference to the hook registry
274    pub fn hook_registry(&self) -> &Arc<HookRegistry<AgentState>> {
275        &self.hook_registry
276    }
277
278    /// Get the model options
279    pub fn model_options(&self) -> &ModelOptions {
280        &self.model_options
281    }
282
283    /// Get reference to the session storage
284    ///
285    /// Use this for all session and checkpoint operations.
286    pub fn session_storage(&self) -> &Arc<dyn SessionStorage> {
287        &self.session_storage
288    }
289}
290
291// Debug implementation for AgentClient
292impl std::fmt::Debug for AgentClient {
293    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
294        f.debug_struct("AgentClient")
295            .field("has_stakpak", &self.has_stakpak())
296            .field("model_options", &self.model_options)
297            .finish_non_exhaustive()
298    }
299}
stakpak_api/client/mod.rs

stakpak_api/client/
mod.rs
