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    /// Build just the `SessionStorage` backend for a given config, without
127    /// initializing LLM providers, hook registries, or the StakpakApiClient.
128    ///
129    /// Use this for read-only session inspection commands where initializing
130    /// a full `AgentClient` would needlessly refresh OAuth tokens and instantiate
131    /// providers the command never uses.
132    pub async fn build_session_storage(
133        stakpak: Option<StakpakConfig>,
134        store_path: Option<String>,
135        profile_name: Option<String>,
136    ) -> Result<Arc<dyn SessionStorage>, String> {
137        if let Some(stakpak) = stakpak
138            && !stakpak.api_key.is_empty()
139        {
140            let storage = StakpakStorage::new_with_profile(
141                &stakpak.api_key,
142                &stakpak.api_endpoint,
143                profile_name,
144            )
145            .map_err(|e| format!("Failed to create Stakpak storage: {}", e))?;
146            Ok(Arc::new(storage))
147        } else {
148            let store_path = store_path.unwrap_or_else(|| {
149                std::env::var("HOME")
150                    .map(|h| format!("{}/{}", h, DEFAULT_STORE_PATH))
151                    .unwrap_or_else(|_| DEFAULT_STORE_PATH.to_string())
152            });
153            let storage = LocalStorage::new(&store_path)
154                .await
155                .map_err(|e| format!("Failed to create local storage: {}", e))?;
156            Ok(Arc::new(storage))
157        }
158    }
159
160    /// Create a new AgentClient
161    pub async fn new(config: AgentClientConfig) -> Result<Self, String> {
162        // 1. Build LLMProviderConfig with Stakpak if configured (only if api_key is not empty)
163        let mut providers = config.providers.clone();
164        if let Some(stakpak) = &config.stakpak
165            && !stakpak.api_key.is_empty()
166        {
167            providers.providers.insert(
168                "stakpak".to_string(),
169                ProviderConfig::Stakpak {
170                    api_key: Some(stakpak.api_key.clone()),
171                    api_endpoint: Some(stakpak.api_endpoint.clone()),
172                    auth: None,
173                },
174            );
175        }
176
177        // 2. Create StakAIClient with all providers
178        let stakai = StakAIClient::new(&providers)
179            .map_err(|e| format!("Failed to create StakAI client: {}", e))?;
180
181        // 3. Create StakpakApiClient if configured (only if api_key is not empty)
182        let stakpak_api = if let Some(stakpak) = &config.stakpak {
183            if !stakpak.api_key.is_empty() {
184                Some(
185                    StakpakApiClient::new(&StakpakApiConfig {
186                        api_key: stakpak.api_key.clone(),
187                        api_endpoint: stakpak.api_endpoint.clone(),
188                    })
189                    .map_err(|e| format!("Failed to create Stakpak API client: {}", e))?,
190                )
191            } else {
192                None
193            }
194        } else {
195            None
196        };
197
198        // 4. Create session storage (Stakpak API or local SQLite)
199        let session_storage: Arc<dyn SessionStorage> = if let Some(stakpak) = &config.stakpak
200            && !stakpak.api_key.is_empty()
201        {
202            Arc::new(
203                StakpakStorage::new(&stakpak.api_key, &stakpak.api_endpoint)
204                    .map_err(|e| format!("Failed to create Stakpak storage: {}", e))?,
205            )
206        } else {
207            let store_path = config.store_path.clone().unwrap_or_else(|| {
208                std::env::var("HOME")
209                    .map(|h| format!("{}/{}", h, DEFAULT_STORE_PATH))
210                    .unwrap_or_else(|_| DEFAULT_STORE_PATH.to_string())
211            });
212            Arc::new(
213                LocalStorage::new(&store_path)
214                    .await
215                    .map_err(|e| format!("Failed to create local storage: {}", e))?,
216            )
217        };
218
219        // 6. Setup hook registry with context management hooks
220        let mut hook_registry = config.hook_registry.unwrap_or_default();
221        hook_registry.register(
222            LifecycleEvent::BeforeInference,
223            Box::new(TaskBoardContextHook::new(TaskBoardContextHookOptions {
224                keep_last_n_assistant_messages: Some(5), // Keep the last 5 assistant messages in context
225                context_budget_threshold: Some(0.8),     // defaults to 0.8 (80%)
226            })),
227        );
228        let hook_registry = Arc::new(hook_registry);
229
230        Ok(Self {
231            stakai,
232            stakpak_api,
233            session_storage,
234            hook_registry,
235            stakpak: config.stakpak,
236        })
237    }
238
239    /// Check if Stakpak API is available
240    pub fn has_stakpak(&self) -> bool {
241        self.stakpak_api.is_some()
242    }
243
244    /// Get the Stakpak API endpoint (with default fallback)
245    pub fn get_stakpak_api_endpoint(&self) -> &str {
246        self.stakpak
247            .as_ref()
248            .map(|s| s.api_endpoint.as_str())
249            .unwrap_or(DEFAULT_STAKPAK_ENDPOINT)
250    }
251
252    /// Get reference to the StakAI client
253    pub fn stakai(&self) -> &StakAIClient {
254        &self.stakai
255    }
256
257    /// Get reference to the Stakpak API client (if available)
258    pub fn stakpak_api(&self) -> Option<&StakpakApiClient> {
259        self.stakpak_api.as_ref()
260    }
261
262    /// Get reference to the hook registry
263    pub fn hook_registry(&self) -> &Arc<HookRegistry<AgentState>> {
264        &self.hook_registry
265    }
266
267    /// Get reference to the session storage
268    ///
269    /// Use this for all session and checkpoint operations.
270    pub fn session_storage(&self) -> &Arc<dyn SessionStorage> {
271        &self.session_storage
272    }
273}
274
275// Debug implementation for AgentClient
276impl std::fmt::Debug for AgentClient {
277    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
278        f.debug_struct("AgentClient")
279            .field("has_stakpak", &self.has_stakpak())
280            .finish_non_exhaustive()
281    }
282}
stakpak_api/client/mod.rs

stakpak_api/client/
mod.rs
