distri_types/
stores.rs

1use crate::{
2    ScratchpadEntry, ToolAuthStore, ToolResponse,
3};
4use async_trait::async_trait;
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize, de::DeserializeOwned};
7use serde_json::Value;
8use std::{collections::HashMap, sync::Arc};
9use tokio::sync::oneshot;
10use uuid::Uuid;
11
12use crate::{
13    AgentEvent, CreateThreadRequest, Message, Task, TaskMessage, TaskStatus, Thread,
14    UpdateThreadRequest,
15};
16
17// Redis and PostgreSQL stores moved to distri-stores crate
18
19/// Filter for listing threads
20#[derive(Debug, Clone, Default, Serialize, Deserialize)]
21pub struct ThreadListFilter {
22    /// Filter by agent ID
23    pub agent_id: Option<String>,
24    /// Filter by external ID (for integration with external systems)
25    pub external_id: Option<String>,
26    /// Filter by thread attributes (JSON matching)
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub attributes: Option<serde_json::Value>,
29    /// Full-text search across title and last_message
30    pub search: Option<String>,
31    /// Filter threads updated after this time
32    pub from_date: Option<DateTime<Utc>>,
33    /// Filter threads updated before this time
34    pub to_date: Option<DateTime<Utc>>,
35    /// Filter by tags (array of tag strings to match)
36    pub tags: Option<Vec<String>>,
37}
38
39/// Paginated response for thread listing
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct ThreadListResponse {
42    pub threads: Vec<crate::ThreadSummary>,
43    pub total: i64,
44    pub page: u32,
45    pub page_size: u32,
46}
47
48/// Agent usage information for sorting agents by thread count
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct AgentUsageInfo {
51    pub agent_id: String,
52    pub agent_name: String,
53    pub thread_count: i64,
54}
55
56/// Initialized store collection
57#[derive(Clone)]
58pub struct InitializedStores {
59    pub session_store: Arc<dyn SessionStore>,
60    pub agent_store: Arc<dyn AgentStore>,
61    pub task_store: Arc<dyn TaskStore>,
62    pub thread_store: Arc<dyn ThreadStore>,
63    pub tool_auth_store: Arc<dyn ToolAuthStore>,
64    pub scratchpad_store: Arc<dyn ScratchpadStore>,
65    pub memory_store: Option<Arc<dyn MemoryStore>>,
66    pub crawl_store: Option<Arc<dyn CrawlStore>>,
67    pub external_tool_calls_store: Arc<dyn ExternalToolCallsStore>,
68    pub prompt_template_store: Option<Arc<dyn PromptTemplateStore>>,
69    pub secret_store: Option<Arc<dyn SecretStore>>,
70    pub skill_store: Option<Arc<dyn SkillStore>>,
71}
72impl InitializedStores {
73    pub fn set_tool_auth_store(&mut self, tool_auth_store: Arc<dyn ToolAuthStore>) {
74        self.tool_auth_store = tool_auth_store;
75    }
76
77    pub fn set_external_tool_calls_store(mut self, store: Arc<dyn ExternalToolCallsStore>) {
78        self.external_tool_calls_store = store;
79    }
80
81    pub fn set_session_store(&mut self, session_store: Arc<dyn SessionStore>) {
82        self.session_store = session_store;
83    }
84
85    pub fn set_agent_store(&mut self, agent_store: Arc<dyn AgentStore>) {
86        self.agent_store = agent_store;
87    }
88
89    pub fn with_task_store(&mut self, task_store: Arc<dyn TaskStore>) {
90        self.task_store = task_store;
91    }
92
93    pub fn with_thread_store(&mut self, thread_store: Arc<dyn ThreadStore>) {
94        self.thread_store = thread_store;
95    }
96
97    pub fn with_scratchpad_store(&mut self, scratchpad_store: Arc<dyn ScratchpadStore>) {
98        self.scratchpad_store = scratchpad_store;
99    }
100}
101
102impl std::fmt::Debug for InitializedStores {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        f.debug_struct("InitializedStores").finish()
105    }
106}
107
108#[derive(Debug, Serialize, Deserialize, Clone)]
109pub struct SessionSummary {
110    pub session_id: String,
111    pub keys: Vec<String>,
112    pub key_count: usize,
113    pub updated_at: Option<DateTime<Utc>>,
114}
115
116// SessionStore trait - manages current conversation thread/run
117#[async_trait::async_trait]
118pub trait SessionStore: Send + Sync + std::fmt::Debug {
119    async fn clear_session(&self, namespace: &str) -> anyhow::Result<()>;
120
121    async fn set_value(&self, namespace: &str, key: &str, value: &Value) -> anyhow::Result<()>;
122
123    async fn set_value_with_expiry(
124        &self,
125        namespace: &str,
126        key: &str,
127        value: &Value,
128        expiry: Option<chrono::DateTime<chrono::Utc>>,
129    ) -> anyhow::Result<()>;
130
131    async fn get_value(&self, namespace: &str, key: &str) -> anyhow::Result<Option<Value>>;
132
133    async fn delete_value(&self, namespace: &str, key: &str) -> anyhow::Result<()>;
134
135    async fn get_all_values(&self, namespace: &str) -> anyhow::Result<HashMap<String, Value>>;
136
137    async fn list_sessions(
138        &self,
139        namespace: Option<&str>,
140        limit: Option<usize>,
141        offset: Option<usize>,
142    ) -> anyhow::Result<Vec<SessionSummary>>;
143}
144#[async_trait::async_trait]
145pub trait SessionStoreExt: SessionStore {
146    async fn set<T: Serialize + Sync>(
147        &self,
148        namespace: &str,
149        key: &str,
150        value: &T,
151    ) -> anyhow::Result<()> {
152        self.set_value(namespace, key, &serde_json::to_value(value)?)
153            .await
154    }
155    async fn set_with_expiry<T: Serialize + Sync>(
156        &self,
157        namespace: &str,
158        key: &str,
159        value: &T,
160        expiry: Option<chrono::DateTime<chrono::Utc>>,
161    ) -> anyhow::Result<()> {
162        self.set_value_with_expiry(namespace, key, &serde_json::to_value(value)?, expiry)
163            .await
164    }
165    async fn get<T: DeserializeOwned + Sync>(
166        &self,
167        namespace: &str,
168        key: &str,
169    ) -> anyhow::Result<Option<T>> {
170        match self.get_value(namespace, key).await? {
171            Some(b) => Ok(Some(serde_json::from_value(b)?)),
172            None => Ok(None),
173        }
174    }
175}
176impl<T: SessionStore + ?Sized> SessionStoreExt for T {}
177
178// Higher-level MemoryStore trait - manages cross-session permanent memory using user_id
179#[async_trait::async_trait]
180pub trait MemoryStore: Send + Sync {
181    /// Store permanent memory from a session for cross-session access
182    async fn store_memory(
183        &self,
184        user_id: &str,
185        session_memory: SessionMemory,
186    ) -> anyhow::Result<()>;
187
188    /// Search for relevant memories across sessions for a user
189    async fn search_memories(
190        &self,
191        user_id: &str,
192        query: &str,
193        limit: Option<usize>,
194    ) -> anyhow::Result<Vec<String>>;
195
196    /// Get all permanent memories for a user
197    async fn get_user_memories(&self, user_id: &str) -> anyhow::Result<Vec<String>>;
198
199    /// Clear all memories for a user
200    async fn clear_user_memories(&self, user_id: &str) -> anyhow::Result<()>;
201}
202
203#[derive(Debug, Clone)]
204pub struct SessionMemory {
205    pub agent_id: String,
206    pub thread_id: String,
207    pub session_summary: String,
208    pub key_insights: Vec<String>,
209    pub important_facts: Vec<String>,
210    pub timestamp: chrono::DateTime<chrono::Utc>,
211}
212#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
213#[serde(tag = "type", rename_all = "snake_case")]
214pub enum FilterMessageType {
215    Events,
216    Messages,
217    Artifacts,
218}
219
220#[derive(Debug, Clone, Serialize, Deserialize)]
221pub struct MessageFilter {
222    pub filter: Option<Vec<FilterMessageType>>,
223    pub limit: Option<usize>,
224    pub offset: Option<usize>,
225}
226
227// Task Store trait for A2A task management
228#[async_trait]
229pub trait TaskStore: Send + Sync {
230    fn init_task(
231        &self,
232        context_id: &str,
233        task_id: Option<&str>,
234        status: Option<TaskStatus>,
235    ) -> Task {
236        let task_id = task_id.unwrap_or(&Uuid::new_v4().to_string()).to_string();
237        Task {
238            id: task_id,
239            status: status.unwrap_or(TaskStatus::Pending),
240            created_at: chrono::Utc::now().timestamp_millis(),
241            updated_at: chrono::Utc::now().timestamp_millis(),
242            thread_id: context_id.to_string(),
243            parent_task_id: None,
244        }
245    }
246    async fn get_or_create_task(
247        &self,
248        thread_id: &str,
249        task_id: &str,
250    ) -> Result<(), anyhow::Error> {
251        match self.get_task(&task_id).await? {
252            Some(task) => task,
253            None => {
254                self.create_task(&thread_id, Some(&task_id), Some(TaskStatus::Running))
255                    .await?
256            }
257        };
258
259        Ok(())
260    }
261    async fn create_task(
262        &self,
263        context_id: &str,
264        task_id: Option<&str>,
265        task_status: Option<TaskStatus>,
266    ) -> anyhow::Result<Task>;
267    async fn get_task(&self, task_id: &str) -> anyhow::Result<Option<Task>>;
268    async fn update_task_status(&self, task_id: &str, status: TaskStatus) -> anyhow::Result<()>;
269    async fn add_event_to_task(&self, task_id: &str, event: AgentEvent) -> anyhow::Result<()>;
270    async fn add_message_to_task(&self, task_id: &str, message: &Message) -> anyhow::Result<()>;
271    async fn cancel_task(&self, task_id: &str) -> anyhow::Result<Task>;
272    async fn list_tasks(&self, thread_id: Option<&str>) -> anyhow::Result<Vec<Task>>;
273
274    async fn get_history(
275        &self,
276        thread_id: &str,
277        filter: Option<MessageFilter>,
278    ) -> anyhow::Result<Vec<(Task, Vec<TaskMessage>)>>;
279
280    async fn update_parent_task(
281        &self,
282        task_id: &str,
283        parent_task_id: Option<&str>,
284    ) -> anyhow::Result<()>;
285}
286
287// Thread Store trait for thread management
288#[async_trait]
289pub trait ThreadStore: Send + Sync {
290    fn as_any(&self) -> &dyn std::any::Any;
291    async fn create_thread(&self, request: CreateThreadRequest) -> anyhow::Result<Thread>;
292    async fn get_thread(&self, thread_id: &str) -> anyhow::Result<Option<Thread>>;
293    async fn update_thread(
294        &self,
295        thread_id: &str,
296        request: UpdateThreadRequest,
297    ) -> anyhow::Result<Thread>;
298    async fn delete_thread(&self, thread_id: &str) -> anyhow::Result<()>;
299
300    /// List threads with pagination and filtering
301    /// Returns a paginated response with total count
302    async fn list_threads(
303        &self,
304        filter: &ThreadListFilter,
305        limit: Option<u32>,
306        offset: Option<u32>,
307    ) -> anyhow::Result<ThreadListResponse>;
308
309    async fn update_thread_with_message(
310        &self,
311        thread_id: &str,
312        message: &str,
313    ) -> anyhow::Result<()>;
314
315    /// Get aggregated home statistics
316    async fn get_home_stats(&self) -> anyhow::Result<HomeStats>;
317
318    /// Get agents sorted by thread count (most active first)
319    /// Includes all registered agents (even those with 0 threads).
320    /// Optionally filters by name using a search string.
321    async fn get_agents_by_usage(&self, search: Option<&str>) -> anyhow::Result<Vec<AgentUsageInfo>>;
322
323    /// Get a map of agent name -> stats for all agents with activity
324    async fn get_agent_stats_map(
325        &self,
326    ) -> anyhow::Result<std::collections::HashMap<String, AgentStatsInfo>>;
327
328    // ========== Message Read Status Methods ==========
329
330    /// Mark a message as read by the current user
331    async fn mark_message_read(
332        &self,
333        thread_id: &str,
334        message_id: &str,
335    ) -> anyhow::Result<MessageReadStatus>;
336
337    /// Get read status for a specific message
338    async fn get_message_read_status(
339        &self,
340        thread_id: &str,
341        message_id: &str,
342    ) -> anyhow::Result<Option<MessageReadStatus>>;
343
344    /// Get read status for all messages in a thread for the current user
345    async fn get_thread_read_status(
346        &self,
347        thread_id: &str,
348    ) -> anyhow::Result<Vec<MessageReadStatus>>;
349
350    // ========== Message Voting Methods ==========
351
352    /// Vote on a message (upvote or downvote)
353    /// For downvotes, a comment is required
354    async fn vote_message(&self, request: VoteMessageRequest) -> anyhow::Result<MessageVote>;
355
356    /// Remove a vote from a message
357    async fn remove_vote(&self, thread_id: &str, message_id: &str) -> anyhow::Result<()>;
358
359    /// Get the current user's vote on a message
360    async fn get_user_vote(
361        &self,
362        thread_id: &str,
363        message_id: &str,
364    ) -> anyhow::Result<Option<MessageVote>>;
365
366    /// Get vote summary for a message (counts + current user's vote)
367    async fn get_message_vote_summary(
368        &self,
369        thread_id: &str,
370        message_id: &str,
371    ) -> anyhow::Result<MessageVoteSummary>;
372
373    /// Get all votes for a message (admin/analytics use)
374    async fn get_message_votes(
375        &self,
376        thread_id: &str,
377        message_id: &str,
378    ) -> anyhow::Result<Vec<MessageVote>>;
379}
380
381/// Home statistics for dashboard
382#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
383pub struct HomeStats {
384    pub total_agents: i64,
385    pub total_threads: i64,
386    pub total_messages: i64,
387    pub avg_run_time_ms: Option<f64>,
388    // Cloud-specific fields (optional)
389    #[serde(skip_serializing_if = "Option::is_none")]
390    pub total_owned_agents: Option<i64>,
391    #[serde(skip_serializing_if = "Option::is_none")]
392    pub total_accessible_agents: Option<i64>,
393    #[serde(skip_serializing_if = "Option::is_none")]
394    pub most_active_agent: Option<MostActiveAgent>,
395    #[serde(skip_serializing_if = "Option::is_none")]
396    pub latest_threads: Option<Vec<LatestThreadInfo>>,
397    /// Recently used agents (last 10 by most recent thread activity)
398    #[serde(skip_serializing_if = "Option::is_none")]
399    pub recently_used_agents: Option<Vec<RecentlyUsedAgent>>,
400    /// Custom metrics that can be displayed in the stats overview
401    /// Key is the metric name (e.g., "usage"), value is the metric data
402    #[serde(skip_serializing_if = "Option::is_none")]
403    pub custom_metrics: Option<std::collections::HashMap<String, CustomMetric>>,
404}
405
406/// A custom metric for display in the stats overview
407#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
408pub struct CustomMetric {
409    /// Display label (e.g., "Monthly Calls")
410    pub label: String,
411    /// Current value as a string (formatted)
412    pub value: String,
413    /// Optional helper text below the value
414    #[serde(skip_serializing_if = "Option::is_none")]
415    pub helper: Option<String>,
416    /// Optional limit (for progress display)
417    #[serde(skip_serializing_if = "Option::is_none")]
418    pub limit: Option<String>,
419    /// Optional raw numeric value for calculations
420    #[serde(skip_serializing_if = "Option::is_none")]
421    pub raw_value: Option<i64>,
422    /// Optional raw limit for calculations
423    #[serde(skip_serializing_if = "Option::is_none")]
424    pub raw_limit: Option<i64>,
425}
426
427#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
428pub struct MostActiveAgent {
429    pub id: String,
430    pub name: String,
431    pub thread_count: i64,
432}
433
434/// Agent that was recently used (based on thread activity)
435#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
436pub struct RecentlyUsedAgent {
437    pub id: String,
438    pub name: String,
439    pub description: Option<String>,
440    pub last_used_at: chrono::DateTime<chrono::Utc>,
441}
442
443#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
444pub struct LatestThreadInfo {
445    pub id: String,
446    pub title: String,
447    pub agent_id: String,
448    pub agent_name: String,
449    pub updated_at: chrono::DateTime<chrono::Utc>,
450}
451
452/// Agent statistics for display
453#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
454pub struct AgentStatsInfo {
455    pub thread_count: i64,
456    pub sub_agent_usage_count: i64,
457    pub last_used_at: Option<chrono::DateTime<chrono::Utc>>,
458}
459
460#[async_trait]
461pub trait AgentStore: Send + Sync {
462    async fn list(
463        &self,
464        cursor: Option<String>,
465        limit: Option<usize>,
466    ) -> (Vec<crate::configuration::AgentConfig>, Option<String>);
467
468    async fn get(&self, name: &str) -> Option<crate::configuration::AgentConfig>;
469    async fn register(&self, config: crate::configuration::AgentConfig) -> anyhow::Result<()>;
470    /// Update an existing agent with new definition
471    async fn update(&self, config: crate::configuration::AgentConfig) -> anyhow::Result<()>;
472
473    async fn clear(&self) -> anyhow::Result<()>;
474}
475
476/// Store for managing scratchpad entries across conversations
477#[async_trait::async_trait]
478pub trait ScratchpadStore: Send + Sync + std::fmt::Debug {
479    /// Add a scratchpad entry for a specific thread
480    async fn add_entry(
481        &self,
482        thread_id: &str,
483        entry: ScratchpadEntry,
484    ) -> Result<(), crate::AgentError>;
485
486    /// Clear all scratchpad entries for a thread
487    async fn clear_entries(&self, thread_id: &str) -> Result<(), crate::AgentError>;
488
489    /// Get entries for a specific task within a thread
490    async fn get_entries(
491        &self,
492        thread_id: &str,
493        task_id: &str,
494        limit: Option<usize>,
495    ) -> Result<Vec<ScratchpadEntry>, crate::AgentError>;
496
497    async fn get_all_entries(
498        &self,
499        thread_id: &str,
500        limit: Option<usize>,
501    ) -> Result<Vec<ScratchpadEntry>, crate::AgentError>;
502}
503
504/// Web crawl result data
505#[derive(Debug, Clone, Serialize, Deserialize)]
506pub struct CrawlResult {
507    pub id: String,
508    pub url: String,
509    pub title: Option<String>,
510    pub content: String,
511    pub html: Option<String>,
512    pub metadata: serde_json::Value,
513    pub links: Vec<String>,
514    pub images: Vec<String>,
515    pub status_code: Option<u16>,
516    pub crawled_at: chrono::DateTime<chrono::Utc>,
517    pub processing_time_ms: Option<u64>,
518}
519
520/// Store for managing web crawl results
521#[async_trait]
522pub trait CrawlStore: Send + Sync {
523    /// Store a crawl result
524    async fn store_crawl_result(&self, result: CrawlResult) -> anyhow::Result<String>;
525
526    /// Get a crawl result by ID
527    async fn get_crawl_result(&self, id: &str) -> anyhow::Result<Option<CrawlResult>>;
528
529    /// Get crawl results for a specific URL
530    async fn get_crawl_results_by_url(&self, url: &str) -> anyhow::Result<Vec<CrawlResult>>;
531
532    /// Get recent crawl results (within time limit)
533    async fn get_recent_crawl_results(
534        &self,
535        limit: Option<usize>,
536        since: Option<chrono::DateTime<chrono::Utc>>,
537    ) -> anyhow::Result<Vec<CrawlResult>>;
538
539    /// Check if URL was recently crawled (within cache duration)
540    async fn is_url_recently_crawled(
541        &self,
542        url: &str,
543        cache_duration: chrono::Duration,
544    ) -> anyhow::Result<Option<CrawlResult>>;
545
546    /// Delete crawl result
547    async fn delete_crawl_result(&self, id: &str) -> anyhow::Result<()>;
548
549    /// Clear all crawl results older than specified date
550    async fn cleanup_old_results(
551        &self,
552        before: chrono::DateTime<chrono::Utc>,
553    ) -> anyhow::Result<usize>;
554}
555
556// ========== Message Read & Voting Types ==========
557
558/// Vote type for message feedback
559#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
560#[serde(rename_all = "lowercase")]
561pub enum VoteType {
562    Upvote,
563    Downvote,
564}
565
566/// Record of a message being read
567#[derive(Debug, Clone, Serialize, Deserialize)]
568pub struct MessageReadStatus {
569    pub thread_id: String,
570    pub message_id: String,
571    pub user_id: String,
572    pub read_at: chrono::DateTime<chrono::Utc>,
573}
574
575/// Request to mark a message as read
576#[derive(Debug, Clone, Serialize, Deserialize)]
577pub struct MarkMessageReadRequest {
578    pub thread_id: String,
579    pub message_id: String,
580}
581
582/// A vote on a message with optional feedback comment
583#[derive(Debug, Clone, Serialize, Deserialize)]
584pub struct MessageVote {
585    pub id: String,
586    pub thread_id: String,
587    pub message_id: String,
588    pub user_id: String,
589    pub vote_type: VoteType,
590    /// Comment is required for downvotes, optional for upvotes
591    pub comment: Option<String>,
592    pub created_at: chrono::DateTime<chrono::Utc>,
593    pub updated_at: chrono::DateTime<chrono::Utc>,
594}
595
596/// Request to vote on a message
597#[derive(Debug, Clone, Serialize, Deserialize)]
598pub struct VoteMessageRequest {
599    pub thread_id: String,
600    pub message_id: String,
601    pub vote_type: VoteType,
602    /// Required for downvotes
603    pub comment: Option<String>,
604}
605
606/// Summary of votes for a message
607#[derive(Debug, Clone, Serialize, Deserialize, Default)]
608pub struct MessageVoteSummary {
609    pub message_id: String,
610    pub upvotes: i64,
611    pub downvotes: i64,
612    /// Current user's vote on this message, if any
613    pub user_vote: Option<VoteType>,
614}
615
616/// Store for managing external tool call completions using oneshot channels
617#[async_trait]
618pub trait ExternalToolCallsStore: Send + Sync + std::fmt::Debug {
619    /// Register a new external tool call session and return a receiver for the response
620    async fn register_external_tool_call(
621        &self,
622        session_id: &str,
623    ) -> anyhow::Result<oneshot::Receiver<ToolResponse>>;
624
625    /// Complete an external tool call by sending the response
626    async fn complete_external_tool_call(
627        &self,
628        session_id: &str,
629        tool_response: ToolResponse,
630    ) -> anyhow::Result<()>;
631
632    /// Remove a session (cleanup)
633    async fn remove_tool_call(&self, session_id: &str) -> anyhow::Result<()>;
634
635    /// List all pending sessions (for debugging)
636    async fn list_pending_tool_calls(&self) -> anyhow::Result<Vec<String>>;
637}
638
639// ========== Prompt Template Store ==========
640
641#[derive(Debug, Clone, Serialize, Deserialize)]
642pub struct PromptTemplateRecord {
643    pub id: String,
644    pub name: String,
645    pub template: String,
646    pub description: Option<String>,
647    pub version: Option<String>,
648    pub is_system: bool,
649    pub created_at: chrono::DateTime<chrono::Utc>,
650    pub updated_at: chrono::DateTime<chrono::Utc>,
651}
652
653#[derive(Debug, Clone, Serialize, Deserialize)]
654pub struct NewPromptTemplate {
655    pub name: String,
656    pub template: String,
657    pub description: Option<String>,
658    pub version: Option<String>,
659    #[serde(default)]
660    pub is_system: bool,
661}
662
663#[derive(Debug, Clone, Serialize, Deserialize)]
664pub struct UpdatePromptTemplate {
665    pub name: String,
666    pub template: String,
667    pub description: Option<String>,
668}
669
670#[async_trait]
671pub trait PromptTemplateStore: Send + Sync {
672    async fn list(&self) -> anyhow::Result<Vec<PromptTemplateRecord>>;
673    async fn get(&self, id: &str) -> anyhow::Result<Option<PromptTemplateRecord>>;
674    async fn create(&self, template: NewPromptTemplate) -> anyhow::Result<PromptTemplateRecord>;
675    async fn update(
676        &self,
677        id: &str,
678        update: UpdatePromptTemplate,
679    ) -> anyhow::Result<PromptTemplateRecord>;
680    async fn delete(&self, id: &str) -> anyhow::Result<()>;
681    async fn clone_template(&self, id: &str) -> anyhow::Result<PromptTemplateRecord>;
682    async fn sync_system_templates(&self, templates: Vec<NewPromptTemplate>) -> anyhow::Result<()>;
683}
684
685// ========== Secret Store ==========
686
687#[derive(Debug, Clone, Serialize, Deserialize)]
688pub struct SecretRecord {
689    pub id: String,
690    pub key: String,
691    pub value: String,
692    pub created_at: chrono::DateTime<chrono::Utc>,
693    pub updated_at: chrono::DateTime<chrono::Utc>,
694}
695
696#[derive(Debug, Clone, Serialize, Deserialize)]
697pub struct NewSecret {
698    pub key: String,
699    pub value: String,
700}
701
702#[async_trait]
703pub trait SecretStore: Send + Sync {
704    async fn list(&self) -> anyhow::Result<Vec<SecretRecord>>;
705    async fn get(&self, key: &str) -> anyhow::Result<Option<SecretRecord>>;
706    async fn create(&self, secret: NewSecret) -> anyhow::Result<SecretRecord>;
707    async fn update(&self, key: &str, value: &str) -> anyhow::Result<SecretRecord>;
708    async fn delete(&self, key: &str) -> anyhow::Result<()>;
709}
710
711// ========== Skill Store ==========
712
713#[derive(Debug, Clone, Serialize, Deserialize)]
714pub struct SkillRecord {
715    pub id: String,
716    pub name: String,
717    pub description: Option<String>,
718    pub content: String,
719    pub tags: Vec<String>,
720    pub is_public: bool,
721    pub is_system: bool,
722    pub star_count: i32,
723    pub clone_count: i32,
724    pub scripts: Vec<SkillScriptRecord>,
725    pub created_at: chrono::DateTime<chrono::Utc>,
726    pub updated_at: chrono::DateTime<chrono::Utc>,
727}
728
729#[derive(Debug, Clone, Serialize, Deserialize)]
730pub struct SkillScriptRecord {
731    pub id: String,
732    pub skill_id: String,
733    pub name: String,
734    pub description: Option<String>,
735    pub code: String,
736    pub language: String,
737    pub created_at: chrono::DateTime<chrono::Utc>,
738    pub updated_at: chrono::DateTime<chrono::Utc>,
739}
740
741#[derive(Debug, Clone, Serialize, Deserialize)]
742pub struct NewSkill {
743    pub name: String,
744    pub description: Option<String>,
745    pub content: String,
746    #[serde(default)]
747    pub tags: Vec<String>,
748    #[serde(default)]
749    pub is_public: bool,
750    #[serde(default)]
751    pub scripts: Vec<NewSkillScript>,
752}
753
754#[derive(Debug, Clone, Serialize, Deserialize)]
755pub struct NewSkillScript {
756    pub name: String,
757    pub description: Option<String>,
758    pub code: String,
759    #[serde(default = "default_script_language")]
760    pub language: String,
761}
762
763fn default_script_language() -> String {
764    "javascript".to_string()
765}
766
767#[derive(Debug, Clone, Serialize, Deserialize)]
768pub struct UpdateSkill {
769    pub name: Option<String>,
770    pub description: Option<String>,
771    pub content: Option<String>,
772    pub tags: Option<Vec<String>>,
773    pub is_public: Option<bool>,
774}
775
776#[derive(Debug, Clone, Serialize, Deserialize)]
777pub struct UpdateSkillScript {
778    pub name: Option<String>,
779    pub description: Option<String>,
780    pub code: Option<String>,
781    pub language: Option<String>,
782}
783
784#[async_trait]
785pub trait SkillStore: Send + Sync {
786    async fn list_skills(&self) -> anyhow::Result<Vec<SkillRecord>>;
787    async fn get_skill(&self, id: &str) -> anyhow::Result<Option<SkillRecord>>;
788    async fn create_skill(&self, skill: NewSkill) -> anyhow::Result<SkillRecord>;
789    async fn update_skill(&self, id: &str, update: UpdateSkill) -> anyhow::Result<SkillRecord>;
790    async fn delete_skill(&self, id: &str) -> anyhow::Result<()>;
791
792    // Script management
793    async fn add_script(&self, skill_id: &str, script: NewSkillScript) -> anyhow::Result<SkillScriptRecord>;
794    async fn update_script(&self, script_id: &str, update: UpdateSkillScript) -> anyhow::Result<SkillScriptRecord>;
795    async fn delete_script(&self, script_id: &str) -> anyhow::Result<()>;
796
797    // Discovery
798    async fn list_public_skills(&self) -> anyhow::Result<Vec<SkillRecord>>;
799    async fn star_skill(&self, skill_id: &str) -> anyhow::Result<()>;
800    async fn unstar_skill(&self, skill_id: &str) -> anyhow::Result<()>;
801    async fn list_starred_skills(&self) -> anyhow::Result<Vec<SkillRecord>>;
802    async fn clone_skill(&self, skill_id: &str) -> anyhow::Result<SkillRecord>;
803}
804
805// ─── Usage Service ──────────────────────────────────────────────────────────
806
807/// Current usage snapshot for a workspace/user.
808#[derive(Debug, Clone, Default, Serialize, Deserialize)]
809pub struct UsageSnapshot {
810    pub day_tokens: i64,
811    pub week_tokens: i64,
812    pub month_tokens: i64,
813}
814
815/// Configured token limits for a workspace.
816#[derive(Debug, Clone, Default, Serialize, Deserialize)]
817pub struct UsageLimits {
818    pub daily_tokens: Option<i64>,
819    pub weekly_tokens: Option<i64>,
820    pub monthly_tokens: Option<i64>,
821}
822
823/// Result of a rate limit check.
824#[derive(Debug, Clone)]
825pub enum UsageCheckResult {
826    Allowed,
827    Denied { reason: String },
828}
829
830/// Trait for usage tracking, rate limiting, and workspace limit management.
831///
832/// OSS: can use a no-op or in-memory implementation.
833/// Cloud: backed by Redis + Postgres with caching.
834#[async_trait]
835pub trait UsageService: Send + Sync {
836    /// Check whether a request should be allowed based on all rate limits.
837    /// Called by middleware before processing a request.
838    /// `is_llm` indicates whether this is an LLM-consuming endpoint.
839    /// `auth_source` is "jwt" or "api_key" for per-source analytics.
840    async fn check_request(
841        &self,
842        workspace_id: &str,
843        user_id: &str,
844        is_llm: bool,
845        auth_source: &str,
846    ) -> UsageCheckResult;
847
848    /// Record token usage after a completed agent run.
849    async fn record_usage(
850        &self,
851        workspace_id: &str,
852        user_id: &str,
853        tokens_used: i64,
854    ) -> anyhow::Result<()>;
855
856    /// Get current usage snapshot for display.
857    async fn get_usage(
858        &self,
859        workspace_id: &str,
860        user_id: &str,
861    ) -> anyhow::Result<UsageSnapshot>;
862
863    /// Get the configured limits for a workspace.
864    async fn get_limits(
865        &self,
866        workspace_id: &str,
867    ) -> anyhow::Result<UsageLimits>;
868}
869
870/// No-op usage service for OSS / development.
871/// Always allows requests, never records anything.
872#[derive(Debug, Clone)]
873pub struct NoOpUsageService;
874
875#[async_trait]
876impl UsageService for NoOpUsageService {
877    async fn check_request(&self, _workspace_id: &str, _user_id: &str, _is_llm: bool, _auth_source: &str) -> UsageCheckResult {
878        UsageCheckResult::Allowed
879    }
880
881    async fn record_usage(&self, _workspace_id: &str, _user_id: &str, _tokens_used: i64) -> anyhow::Result<()> {
882        Ok(())
883    }
884
885    async fn get_usage(&self, _workspace_id: &str, _user_id: &str) -> anyhow::Result<UsageSnapshot> {
886        Ok(UsageSnapshot::default())
887    }
888
889    async fn get_limits(&self, _workspace_id: &str) -> anyhow::Result<UsageLimits> {
890        Ok(UsageLimits::default())
891    }
892}
distri_types/stores.rs

distri_types/
stores.rs
