distri_types/
stores.rs

1use crate::{ScratchpadEntry, ToolAuthStore, ToolResponse};
2use async_trait::async_trait;
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize, de::DeserializeOwned};
5use serde_json::Value;
6use std::{collections::HashMap, sync::Arc};
7use tokio::sync::oneshot;
8use uuid::Uuid;
9
10use crate::{
11    AgentEvent, CreateThreadRequest, Message, Task, TaskMessage, TaskStatus, Thread,
12    UpdateThreadRequest,
13};
14
15// Redis and PostgreSQL stores moved to distri-stores crate
16
17/// Filter for listing threads
18#[derive(Debug, Clone, Default, Serialize, Deserialize)]
19pub struct ThreadListFilter {
20    /// Filter by agent ID
21    pub agent_id: Option<String>,
22    /// Filter by external ID (for integration with external systems)
23    pub external_id: Option<String>,
24    /// Filter by thread attributes (JSON matching)
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub attributes: Option<serde_json::Value>,
27    /// Full-text search across title and last_message
28    pub search: Option<String>,
29    /// Filter threads updated after this time
30    pub from_date: Option<DateTime<Utc>>,
31    /// Filter threads updated before this time
32    pub to_date: Option<DateTime<Utc>>,
33    /// Filter by tags (array of tag strings to match)
34    pub tags: Option<Vec<String>>,
35}
36
37/// Paginated response for thread listing
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct ThreadListResponse {
40    pub threads: Vec<crate::ThreadSummary>,
41    pub total: i64,
42    pub page: u32,
43    pub page_size: u32,
44}
45
46/// Agent usage information for sorting agents by thread count
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct AgentUsageInfo {
49    pub agent_id: String,
50    pub agent_name: String,
51    pub thread_count: i64,
52}
53
54/// Initialized store collection
55#[derive(Clone)]
56pub struct InitializedStores {
57    pub session_store: Arc<dyn SessionStore>,
58    pub agent_store: Arc<dyn AgentStore>,
59    pub task_store: Arc<dyn TaskStore>,
60    pub thread_store: Arc<dyn ThreadStore>,
61    pub tool_auth_store: Arc<dyn ToolAuthStore>,
62    pub scratchpad_store: Arc<dyn ScratchpadStore>,
63    pub memory_store: Option<Arc<dyn MemoryStore>>,
64    pub crawl_store: Option<Arc<dyn CrawlStore>>,
65    pub external_tool_calls_store: Arc<dyn ExternalToolCallsStore>,
66    pub prompt_template_store: Option<Arc<dyn PromptTemplateStore>>,
67    pub secret_store: Option<Arc<dyn SecretStore>>,
68    pub skill_store: Option<Arc<dyn SkillStore>>,
69    pub workflow_store: Option<Arc<dyn WorkflowStore>>,
70}
71impl InitializedStores {
72    pub fn set_tool_auth_store(&mut self, tool_auth_store: Arc<dyn ToolAuthStore>) {
73        self.tool_auth_store = tool_auth_store;
74    }
75
76    pub fn set_external_tool_calls_store(mut self, store: Arc<dyn ExternalToolCallsStore>) {
77        self.external_tool_calls_store = store;
78    }
79
80    pub fn set_session_store(&mut self, session_store: Arc<dyn SessionStore>) {
81        self.session_store = session_store;
82    }
83
84    pub fn set_agent_store(&mut self, agent_store: Arc<dyn AgentStore>) {
85        self.agent_store = agent_store;
86    }
87
88    pub fn with_task_store(&mut self, task_store: Arc<dyn TaskStore>) {
89        self.task_store = task_store;
90    }
91
92    pub fn with_thread_store(&mut self, thread_store: Arc<dyn ThreadStore>) {
93        self.thread_store = thread_store;
94    }
95
96    pub fn with_scratchpad_store(&mut self, scratchpad_store: Arc<dyn ScratchpadStore>) {
97        self.scratchpad_store = scratchpad_store;
98    }
99}
100
101impl std::fmt::Debug for InitializedStores {
102    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103        f.debug_struct("InitializedStores").finish()
104    }
105}
106
107#[derive(Debug, Serialize, Deserialize, Clone)]
108pub struct SessionSummary {
109    pub session_id: String,
110    pub keys: Vec<String>,
111    pub key_count: usize,
112    pub updated_at: Option<DateTime<Utc>>,
113}
114
115// SessionStore trait - manages current conversation thread/run
116#[async_trait::async_trait]
117pub trait SessionStore: Send + Sync + std::fmt::Debug {
118    async fn clear_session(&self, namespace: &str) -> anyhow::Result<()>;
119
120    async fn set_value(&self, namespace: &str, key: &str, value: &Value) -> anyhow::Result<()>;
121
122    async fn set_value_with_expiry(
123        &self,
124        namespace: &str,
125        key: &str,
126        value: &Value,
127        expiry: Option<chrono::DateTime<chrono::Utc>>,
128    ) -> anyhow::Result<()>;
129
130    async fn get_value(&self, namespace: &str, key: &str) -> anyhow::Result<Option<Value>>;
131
132    async fn delete_value(&self, namespace: &str, key: &str) -> anyhow::Result<()>;
133
134    async fn get_all_values(&self, namespace: &str) -> anyhow::Result<HashMap<String, Value>>;
135
136    async fn list_sessions(
137        &self,
138        namespace: Option<&str>,
139        limit: Option<usize>,
140        offset: Option<usize>,
141    ) -> anyhow::Result<Vec<SessionSummary>>;
142}
143#[async_trait::async_trait]
144pub trait SessionStoreExt: SessionStore {
145    async fn set<T: Serialize + Sync>(
146        &self,
147        namespace: &str,
148        key: &str,
149        value: &T,
150    ) -> anyhow::Result<()> {
151        self.set_value(namespace, key, &serde_json::to_value(value)?)
152            .await
153    }
154    async fn set_with_expiry<T: Serialize + Sync>(
155        &self,
156        namespace: &str,
157        key: &str,
158        value: &T,
159        expiry: Option<chrono::DateTime<chrono::Utc>>,
160    ) -> anyhow::Result<()> {
161        self.set_value_with_expiry(namespace, key, &serde_json::to_value(value)?, expiry)
162            .await
163    }
164    async fn get<T: DeserializeOwned + Sync>(
165        &self,
166        namespace: &str,
167        key: &str,
168    ) -> anyhow::Result<Option<T>> {
169        match self.get_value(namespace, key).await? {
170            Some(b) => Ok(Some(serde_json::from_value(b)?)),
171            None => Ok(None),
172        }
173    }
174}
175impl<T: SessionStore + ?Sized> SessionStoreExt for T {}
176
177// Higher-level MemoryStore trait - manages cross-session permanent memory using user_id
178#[async_trait::async_trait]
179pub trait MemoryStore: Send + Sync {
180    /// Store permanent memory from a session for cross-session access
181    async fn store_memory(
182        &self,
183        user_id: &str,
184        session_memory: SessionMemory,
185    ) -> anyhow::Result<()>;
186
187    /// Search for relevant memories across sessions for a user
188    async fn search_memories(
189        &self,
190        user_id: &str,
191        query: &str,
192        limit: Option<usize>,
193    ) -> anyhow::Result<Vec<String>>;
194
195    /// Get all permanent memories for a user
196    async fn get_user_memories(&self, user_id: &str) -> anyhow::Result<Vec<String>>;
197
198    /// Clear all memories for a user
199    async fn clear_user_memories(&self, user_id: &str) -> anyhow::Result<()>;
200}
201
202#[derive(Debug, Clone)]
203pub struct SessionMemory {
204    pub agent_id: String,
205    pub thread_id: String,
206    pub session_summary: String,
207    pub key_insights: Vec<String>,
208    pub important_facts: Vec<String>,
209    pub timestamp: chrono::DateTime<chrono::Utc>,
210}
211#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
212#[serde(tag = "type", rename_all = "snake_case")]
213pub enum FilterMessageType {
214    Events,
215    Messages,
216    Artifacts,
217}
218
219#[derive(Debug, Clone, Serialize, Deserialize)]
220pub struct MessageFilter {
221    pub filter: Option<Vec<FilterMessageType>>,
222    pub limit: Option<usize>,
223    pub offset: Option<usize>,
224}
225
226// Task Store trait for A2A task management
227#[async_trait]
228pub trait TaskStore: Send + Sync {
229    fn init_task(
230        &self,
231        context_id: &str,
232        task_id: Option<&str>,
233        status: Option<TaskStatus>,
234    ) -> Task {
235        let task_id = task_id.unwrap_or(&Uuid::new_v4().to_string()).to_string();
236        Task {
237            id: task_id,
238            status: status.unwrap_or(TaskStatus::Pending),
239            created_at: chrono::Utc::now().timestamp_millis(),
240            updated_at: chrono::Utc::now().timestamp_millis(),
241            thread_id: context_id.to_string(),
242            parent_task_id: None,
243        }
244    }
245    async fn get_or_create_task(
246        &self,
247        thread_id: &str,
248        task_id: &str,
249    ) -> Result<(), anyhow::Error> {
250        match self.get_task(task_id).await? {
251            Some(task) => task,
252            None => {
253                self.create_task(thread_id, Some(task_id), Some(TaskStatus::Running))
254                    .await?
255            }
256        };
257
258        Ok(())
259    }
260    async fn create_task(
261        &self,
262        context_id: &str,
263        task_id: Option<&str>,
264        task_status: Option<TaskStatus>,
265    ) -> anyhow::Result<Task>;
266    async fn get_task(&self, task_id: &str) -> anyhow::Result<Option<Task>>;
267    async fn update_task_status(&self, task_id: &str, status: TaskStatus) -> anyhow::Result<()>;
268    async fn add_event_to_task(&self, task_id: &str, event: AgentEvent) -> anyhow::Result<()>;
269    async fn add_message_to_task(&self, task_id: &str, message: &Message) -> anyhow::Result<()>;
270    async fn cancel_task(&self, task_id: &str) -> anyhow::Result<Task>;
271    async fn list_tasks(&self, thread_id: Option<&str>) -> anyhow::Result<Vec<Task>>;
272
273    async fn get_history(
274        &self,
275        thread_id: &str,
276        filter: Option<MessageFilter>,
277    ) -> anyhow::Result<Vec<(Task, Vec<TaskMessage>)>>;
278
279    async fn update_parent_task(
280        &self,
281        task_id: &str,
282        parent_task_id: Option<&str>,
283    ) -> anyhow::Result<()>;
284}
285
286// Thread Store trait for thread management
287#[async_trait]
288pub trait ThreadStore: Send + Sync {
289    fn as_any(&self) -> &dyn std::any::Any;
290    async fn create_thread(&self, request: CreateThreadRequest) -> anyhow::Result<Thread>;
291    async fn get_thread(&self, thread_id: &str) -> anyhow::Result<Option<Thread>>;
292    async fn update_thread(
293        &self,
294        thread_id: &str,
295        request: UpdateThreadRequest,
296    ) -> anyhow::Result<Thread>;
297    async fn delete_thread(&self, thread_id: &str) -> anyhow::Result<()>;
298
299    /// List threads with pagination and filtering
300    /// Returns a paginated response with total count
301    async fn list_threads(
302        &self,
303        filter: &ThreadListFilter,
304        limit: Option<u32>,
305        offset: Option<u32>,
306    ) -> anyhow::Result<ThreadListResponse>;
307
308    async fn update_thread_with_message(
309        &self,
310        thread_id: &str,
311        message: &str,
312    ) -> anyhow::Result<()>;
313
314    /// Get aggregated home statistics
315    async fn get_home_stats(&self) -> anyhow::Result<HomeStats>;
316
317    /// Get agents sorted by thread count (most active first)
318    /// Includes all registered agents (even those with 0 threads).
319    /// Optionally filters by name using a search string.
320    async fn get_agents_by_usage(
321        &self,
322        search: Option<&str>,
323    ) -> anyhow::Result<Vec<AgentUsageInfo>>;
324
325    /// Get a map of agent name -> stats for all agents with activity
326    async fn get_agent_stats_map(
327        &self,
328    ) -> anyhow::Result<std::collections::HashMap<String, AgentStatsInfo>>;
329
330    // ========== Message Read Status Methods ==========
331
332    /// Mark a message as read by the current user
333    async fn mark_message_read(
334        &self,
335        thread_id: &str,
336        message_id: &str,
337    ) -> anyhow::Result<MessageReadStatus>;
338
339    /// Get read status for a specific message
340    async fn get_message_read_status(
341        &self,
342        thread_id: &str,
343        message_id: &str,
344    ) -> anyhow::Result<Option<MessageReadStatus>>;
345
346    /// Get read status for all messages in a thread for the current user
347    async fn get_thread_read_status(
348        &self,
349        thread_id: &str,
350    ) -> anyhow::Result<Vec<MessageReadStatus>>;
351
352    // ========== Message Voting Methods ==========
353
354    /// Vote on a message (upvote or downvote)
355    /// For downvotes, a comment is required
356    async fn vote_message(&self, request: VoteMessageRequest) -> anyhow::Result<MessageVote>;
357
358    /// Remove a vote from a message
359    async fn remove_vote(&self, thread_id: &str, message_id: &str) -> anyhow::Result<()>;
360
361    /// Get the current user's vote on a message
362    async fn get_user_vote(
363        &self,
364        thread_id: &str,
365        message_id: &str,
366    ) -> anyhow::Result<Option<MessageVote>>;
367
368    /// Get vote summary for a message (counts + current user's vote)
369    async fn get_message_vote_summary(
370        &self,
371        thread_id: &str,
372        message_id: &str,
373    ) -> anyhow::Result<MessageVoteSummary>;
374
375    /// Get all votes for a message (admin/analytics use)
376    async fn get_message_votes(
377        &self,
378        thread_id: &str,
379        message_id: &str,
380    ) -> anyhow::Result<Vec<MessageVote>>;
381}
382
383/// Home statistics for dashboard
384#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
385pub struct HomeStats {
386    pub total_agents: i64,
387    pub total_threads: i64,
388    pub total_messages: i64,
389    pub avg_run_time_ms: Option<f64>,
390    // Cloud-specific fields (optional)
391    #[serde(skip_serializing_if = "Option::is_none")]
392    pub total_owned_agents: Option<i64>,
393    #[serde(skip_serializing_if = "Option::is_none")]
394    pub total_accessible_agents: Option<i64>,
395    #[serde(skip_serializing_if = "Option::is_none")]
396    pub most_active_agent: Option<MostActiveAgent>,
397    #[serde(skip_serializing_if = "Option::is_none")]
398    pub latest_threads: Option<Vec<LatestThreadInfo>>,
399    /// Recently used agents (last 10 by most recent thread activity)
400    #[serde(skip_serializing_if = "Option::is_none")]
401    pub recently_used_agents: Option<Vec<RecentlyUsedAgent>>,
402    /// Custom metrics that can be displayed in the stats overview
403    /// Key is the metric name (e.g., "usage"), value is the metric data
404    #[serde(skip_serializing_if = "Option::is_none")]
405    pub custom_metrics: Option<std::collections::HashMap<String, CustomMetric>>,
406}
407
408/// A custom metric for display in the stats overview
409#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
410pub struct CustomMetric {
411    /// Display label (e.g., "Monthly Calls")
412    pub label: String,
413    /// Current value as a string (formatted)
414    pub value: String,
415    /// Optional helper text below the value
416    #[serde(skip_serializing_if = "Option::is_none")]
417    pub helper: Option<String>,
418    /// Optional limit (for progress display)
419    #[serde(skip_serializing_if = "Option::is_none")]
420    pub limit: Option<String>,
421    /// Optional raw numeric value for calculations
422    #[serde(skip_serializing_if = "Option::is_none")]
423    pub raw_value: Option<i64>,
424    /// Optional raw limit for calculations
425    #[serde(skip_serializing_if = "Option::is_none")]
426    pub raw_limit: Option<i64>,
427}
428
429#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
430pub struct MostActiveAgent {
431    pub id: String,
432    pub name: String,
433    pub thread_count: i64,
434}
435
436/// Agent that was recently used (based on thread activity)
437#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
438pub struct RecentlyUsedAgent {
439    pub id: String,
440    pub name: String,
441    pub description: Option<String>,
442    pub last_used_at: chrono::DateTime<chrono::Utc>,
443}
444
445#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
446pub struct LatestThreadInfo {
447    pub id: String,
448    pub title: String,
449    pub agent_id: String,
450    pub agent_name: String,
451    pub updated_at: chrono::DateTime<chrono::Utc>,
452}
453
454/// Agent statistics for display
455#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
456pub struct AgentStatsInfo {
457    pub thread_count: i64,
458    pub sub_agent_usage_count: i64,
459    pub last_used_at: Option<chrono::DateTime<chrono::Utc>>,
460}
461
462#[async_trait]
463pub trait AgentStore: Send + Sync {
464    async fn list(
465        &self,
466        cursor: Option<String>,
467        limit: Option<usize>,
468    ) -> (Vec<crate::configuration::AgentConfig>, Option<String>);
469
470    async fn get(&self, name: &str) -> Option<crate::configuration::AgentConfig>;
471    async fn register(&self, config: crate::configuration::AgentConfig) -> anyhow::Result<()>;
472    /// Update an existing agent with new definition
473    async fn update(&self, config: crate::configuration::AgentConfig) -> anyhow::Result<()>;
474
475    async fn clear(&self) -> anyhow::Result<()>;
476
477    /// Delete an agent by name or ID
478    async fn delete(&self, id: &str) -> anyhow::Result<()>;
479}
480
481/// Store for managing scratchpad entries across conversations
482#[async_trait::async_trait]
483pub trait ScratchpadStore: Send + Sync + std::fmt::Debug {
484    /// Add a scratchpad entry for a specific thread
485    async fn add_entry(
486        &self,
487        thread_id: &str,
488        entry: ScratchpadEntry,
489    ) -> Result<(), crate::AgentError>;
490
491    /// Clear all scratchpad entries for a thread
492    async fn clear_entries(&self, thread_id: &str) -> Result<(), crate::AgentError>;
493
494    /// Get entries for a specific task within a thread
495    async fn get_entries(
496        &self,
497        thread_id: &str,
498        task_id: &str,
499        limit: Option<usize>,
500    ) -> Result<Vec<ScratchpadEntry>, crate::AgentError>;
501
502    async fn get_all_entries(
503        &self,
504        thread_id: &str,
505        limit: Option<usize>,
506    ) -> Result<Vec<ScratchpadEntry>, crate::AgentError>;
507}
508
509/// Web crawl result data
510#[derive(Debug, Clone, Serialize, Deserialize)]
511pub struct CrawlResult {
512    pub id: String,
513    pub url: String,
514    pub title: Option<String>,
515    pub content: String,
516    pub html: Option<String>,
517    pub metadata: serde_json::Value,
518    pub links: Vec<String>,
519    pub images: Vec<String>,
520    pub status_code: Option<u16>,
521    pub crawled_at: chrono::DateTime<chrono::Utc>,
522    pub processing_time_ms: Option<u64>,
523}
524
525/// Store for managing web crawl results
526#[async_trait]
527pub trait CrawlStore: Send + Sync {
528    /// Store a crawl result
529    async fn store_crawl_result(&self, result: CrawlResult) -> anyhow::Result<String>;
530
531    /// Get a crawl result by ID
532    async fn get_crawl_result(&self, id: &str) -> anyhow::Result<Option<CrawlResult>>;
533
534    /// Get crawl results for a specific URL
535    async fn get_crawl_results_by_url(&self, url: &str) -> anyhow::Result<Vec<CrawlResult>>;
536
537    /// Get recent crawl results (within time limit)
538    async fn get_recent_crawl_results(
539        &self,
540        limit: Option<usize>,
541        since: Option<chrono::DateTime<chrono::Utc>>,
542    ) -> anyhow::Result<Vec<CrawlResult>>;
543
544    /// Check if URL was recently crawled (within cache duration)
545    async fn is_url_recently_crawled(
546        &self,
547        url: &str,
548        cache_duration: chrono::Duration,
549    ) -> anyhow::Result<Option<CrawlResult>>;
550
551    /// Delete crawl result
552    async fn delete_crawl_result(&self, id: &str) -> anyhow::Result<()>;
553
554    /// Clear all crawl results older than specified date
555    async fn cleanup_old_results(
556        &self,
557        before: chrono::DateTime<chrono::Utc>,
558    ) -> anyhow::Result<usize>;
559}
560
561// ========== Message Read & Voting Types ==========
562
563/// Vote type for message feedback
564#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
565#[serde(rename_all = "lowercase")]
566pub enum VoteType {
567    Upvote,
568    Downvote,
569}
570
571/// Record of a message being read
572#[derive(Debug, Clone, Serialize, Deserialize)]
573pub struct MessageReadStatus {
574    pub thread_id: String,
575    pub message_id: String,
576    pub user_id: String,
577    pub read_at: chrono::DateTime<chrono::Utc>,
578}
579
580/// Request to mark a message as read
581#[derive(Debug, Clone, Serialize, Deserialize)]
582pub struct MarkMessageReadRequest {
583    pub thread_id: String,
584    pub message_id: String,
585}
586
587/// A vote on a message with optional feedback comment
588#[derive(Debug, Clone, Serialize, Deserialize)]
589pub struct MessageVote {
590    pub id: String,
591    pub thread_id: String,
592    pub message_id: String,
593    pub user_id: String,
594    pub vote_type: VoteType,
595    /// Comment is required for downvotes, optional for upvotes
596    pub comment: Option<String>,
597    pub created_at: chrono::DateTime<chrono::Utc>,
598    pub updated_at: chrono::DateTime<chrono::Utc>,
599}
600
601/// Request to vote on a message
602#[derive(Debug, Clone, Serialize, Deserialize)]
603pub struct VoteMessageRequest {
604    pub thread_id: String,
605    pub message_id: String,
606    pub vote_type: VoteType,
607    /// Required for downvotes
608    pub comment: Option<String>,
609}
610
611/// Summary of votes for a message
612#[derive(Debug, Clone, Serialize, Deserialize, Default)]
613pub struct MessageVoteSummary {
614    pub message_id: String,
615    pub upvotes: i64,
616    pub downvotes: i64,
617    /// Current user's vote on this message, if any
618    pub user_vote: Option<VoteType>,
619}
620
621/// Store for managing external tool call completions using oneshot channels
622#[async_trait]
623pub trait ExternalToolCallsStore: Send + Sync + std::fmt::Debug {
624    /// Register a new external tool call session and return a receiver for the response
625    async fn register_external_tool_call(
626        &self,
627        session_id: &str,
628    ) -> anyhow::Result<oneshot::Receiver<ToolResponse>>;
629
630    /// Complete an external tool call by sending the response
631    async fn complete_external_tool_call(
632        &self,
633        session_id: &str,
634        tool_response: ToolResponse,
635    ) -> anyhow::Result<()>;
636
637    /// Remove a session (cleanup)
638    async fn remove_tool_call(&self, session_id: &str) -> anyhow::Result<()>;
639
640    /// List all pending sessions (for debugging)
641    async fn list_pending_tool_calls(&self) -> anyhow::Result<Vec<String>>;
642}
643
644// ========== Prompt Template Store ==========
645
646#[derive(Debug, Clone, Serialize, Deserialize)]
647pub struct PromptTemplateRecord {
648    pub id: String,
649    pub name: String,
650    pub template: String,
651    pub description: Option<String>,
652    pub version: Option<String>,
653    pub is_system: bool,
654    pub created_at: chrono::DateTime<chrono::Utc>,
655    pub updated_at: chrono::DateTime<chrono::Utc>,
656}
657
658#[derive(Debug, Clone, Serialize, Deserialize)]
659pub struct NewPromptTemplate {
660    pub name: String,
661    pub template: String,
662    pub description: Option<String>,
663    pub version: Option<String>,
664    #[serde(default)]
665    pub is_system: bool,
666}
667
668#[derive(Debug, Clone, Serialize, Deserialize)]
669pub struct UpdatePromptTemplate {
670    pub name: String,
671    pub template: String,
672    pub description: Option<String>,
673}
674
675#[async_trait]
676pub trait PromptTemplateStore: Send + Sync {
677    async fn list(&self) -> anyhow::Result<Vec<PromptTemplateRecord>>;
678    async fn get(&self, id: &str) -> anyhow::Result<Option<PromptTemplateRecord>>;
679    async fn create(&self, template: NewPromptTemplate) -> anyhow::Result<PromptTemplateRecord>;
680    async fn update(
681        &self,
682        id: &str,
683        update: UpdatePromptTemplate,
684    ) -> anyhow::Result<PromptTemplateRecord>;
685    async fn delete(&self, id: &str) -> anyhow::Result<()>;
686    async fn clone_template(&self, id: &str) -> anyhow::Result<PromptTemplateRecord>;
687    async fn sync_system_templates(&self, templates: Vec<NewPromptTemplate>) -> anyhow::Result<()>;
688}
689
690// ========== Secret Store ==========
691
692#[derive(Debug, Clone, Serialize, Deserialize)]
693pub struct SecretRecord {
694    pub id: String,
695    pub key: String,
696    pub value: String,
697    pub created_at: chrono::DateTime<chrono::Utc>,
698    pub updated_at: chrono::DateTime<chrono::Utc>,
699}
700
701#[derive(Debug, Clone, Serialize, Deserialize)]
702pub struct NewSecret {
703    pub key: String,
704    pub value: String,
705}
706
707#[async_trait]
708pub trait SecretStore: Send + Sync {
709    async fn list(&self) -> anyhow::Result<Vec<SecretRecord>>;
710    async fn get(&self, key: &str) -> anyhow::Result<Option<SecretRecord>>;
711    async fn create(&self, secret: NewSecret) -> anyhow::Result<SecretRecord>;
712    async fn update(&self, key: &str, value: &str) -> anyhow::Result<SecretRecord>;
713    async fn delete(&self, key: &str) -> anyhow::Result<()>;
714}
715
716// ========== Provider Store ==========
717
718#[derive(Debug, Clone, Serialize, Deserialize)]
719pub struct CustomProviderConfig {
720    pub id: String,
721    pub name: String,
722    pub base_url: String,
723    #[serde(default, skip_serializing_if = "Option::is_none")]
724    pub project_id: Option<String>,
725}
726
727#[derive(Debug, Clone, Serialize, Deserialize)]
728pub struct CustomModelEntry {
729    pub provider: String,
730    pub model: String,
731}
732
733/// A custom connection provider (OAuth integration) stored in workspace settings.
734#[derive(Debug, Clone, Serialize, Deserialize)]
735pub struct ConnectionProviderConfig {
736    /// Unique identifier (e.g., "linear", "figma", "custom_crm")
737    pub id: String,
738    /// Display name
739    pub name: String,
740    /// OAuth2 authorization URL
741    pub authorization_url: String,
742    /// OAuth2 token URL
743    pub token_url: String,
744    /// Optional refresh URL (defaults to token_url)
745    #[serde(default, skip_serializing_if = "Option::is_none")]
746    pub refresh_url: Option<String>,
747    /// Scopes the provider supports
748    #[serde(default)]
749    pub scopes_supported: Vec<String>,
750    /// Default scopes to request
751    #[serde(default)]
752    pub default_scopes: Vec<String>,
753    /// Friendly scope name → full scope string mappings
754    #[serde(default)]
755    pub scope_mappings: std::collections::HashMap<String, String>,
756}
757
758/// Request payload for upserting a provider configuration.
759#[derive(Debug, Clone, Serialize, Deserialize)]
760pub struct UpsertProviderRequest {
761    pub provider_id: String,
762    #[serde(default)]
763    pub secrets: std::collections::HashMap<String, String>,
764    #[serde(default)]
765    pub config: Option<CustomProviderConfig>,
766    #[serde(default)]
767    pub custom_models: Option<Vec<CustomModelEntry>>,
768    /// Default model in "provider/model" format. Empty string or null to clear.
769    #[serde(default)]
770    pub default_model: Option<String>,
771    /// Connection provider config (OAuth integration) to add/update.
772    #[serde(default)]
773    pub connection_provider: Option<ConnectionProviderConfig>,
774}
775
776/// Response after upserting a provider.
777#[derive(Debug, Clone, Serialize, Deserialize)]
778pub struct UpsertProviderResponse {
779    pub provider_id: String,
780    pub secrets_saved: usize,
781    pub config_saved: bool,
782}
783
784#[async_trait]
785pub trait ProviderStore: Send + Sync {
786    async fn upsert_provider(
787        &self,
788        req: UpsertProviderRequest,
789    ) -> anyhow::Result<UpsertProviderResponse>;
790
791    async fn delete_provider(&self, provider_id: &str) -> anyhow::Result<()>;
792
793    async fn get_default_model(&self) -> anyhow::Result<Option<String>>;
794}
795
796// ========== Skill Store ==========
797
798/// API response wrapper for skill list endpoints.
799#[derive(Debug, Clone, Serialize, Deserialize)]
800pub struct SkillsListResponse {
801    pub skills: Vec<SkillListItem>,
802}
803
804/// Lighter skill record for list endpoints — no content or scripts.
805/// Used by both distri-server (OSS) and distri-cloud.
806#[derive(Debug, Clone, Serialize, Deserialize)]
807pub struct SkillListItem {
808    pub id: String,
809    #[serde(default)]
810    pub workspace_slug: String,
811    pub name: String,
812    #[serde(default)]
813    pub full_name: String,
814    #[serde(default)]
815    pub description: Option<String>,
816    #[serde(default)]
817    pub tags: Vec<String>,
818    #[serde(default)]
819    pub is_public: bool,
820    #[serde(default)]
821    pub is_system: bool,
822    #[serde(default)]
823    pub is_owner: bool,
824    #[serde(default)]
825    pub star_count: i32,
826    #[serde(default)]
827    pub clone_count: i32,
828    #[serde(default)]
829    pub is_starred: bool,
830    pub created_at: chrono::DateTime<chrono::Utc>,
831    pub updated_at: chrono::DateTime<chrono::Utc>,
832}
833
834#[derive(Debug, Clone, Serialize, Deserialize)]
835pub struct SkillRecord {
836    pub id: String,
837    pub name: String,
838    pub description: Option<String>,
839    pub content: String,
840    pub tags: Vec<String>,
841    pub is_public: bool,
842    pub is_system: bool,
843    pub star_count: i32,
844    pub clone_count: i32,
845    pub scripts: Vec<SkillScriptRecord>,
846    pub created_at: chrono::DateTime<chrono::Utc>,
847    pub updated_at: chrono::DateTime<chrono::Utc>,
848}
849
850#[derive(Debug, Clone, Serialize, Deserialize)]
851pub struct SkillScriptRecord {
852    pub id: String,
853    pub skill_id: String,
854    pub name: String,
855    pub description: Option<String>,
856    pub code: String,
857    pub language: String,
858    pub created_at: chrono::DateTime<chrono::Utc>,
859    pub updated_at: chrono::DateTime<chrono::Utc>,
860}
861
862#[derive(Debug, Clone, Serialize, Deserialize)]
863pub struct NewSkill {
864    pub name: String,
865    pub description: Option<String>,
866    pub content: String,
867    #[serde(default)]
868    pub tags: Vec<String>,
869    #[serde(default)]
870    pub is_public: bool,
871    #[serde(default)]
872    pub scripts: Vec<NewSkillScript>,
873}
874
875#[derive(Debug, Clone, Serialize, Deserialize)]
876pub struct NewSkillScript {
877    pub name: String,
878    pub description: Option<String>,
879    pub code: String,
880    #[serde(default = "default_script_language")]
881    pub language: String,
882}
883
884fn default_script_language() -> String {
885    "javascript".to_string()
886}
887
888#[derive(Debug, Clone, Serialize, Deserialize)]
889pub struct UpdateSkill {
890    pub name: Option<String>,
891    pub description: Option<String>,
892    pub content: Option<String>,
893    pub tags: Option<Vec<String>>,
894    pub is_public: Option<bool>,
895}
896
897#[derive(Debug, Clone, Serialize, Deserialize)]
898pub struct UpdateSkillScript {
899    pub name: Option<String>,
900    pub description: Option<String>,
901    pub code: Option<String>,
902    pub language: Option<String>,
903}
904
905#[async_trait]
906pub trait SkillStore: Send + Sync {
907    async fn list_skills(&self) -> anyhow::Result<Vec<SkillRecord>>;
908    async fn get_skill(&self, id: &str) -> anyhow::Result<Option<SkillRecord>>;
909    async fn create_skill(&self, skill: NewSkill) -> anyhow::Result<SkillRecord>;
910    async fn update_skill(&self, id: &str, update: UpdateSkill) -> anyhow::Result<SkillRecord>;
911    async fn delete_skill(&self, id: &str) -> anyhow::Result<()>;
912
913    // Script management
914    async fn add_script(
915        &self,
916        skill_id: &str,
917        script: NewSkillScript,
918    ) -> anyhow::Result<SkillScriptRecord>;
919    async fn update_script(
920        &self,
921        script_id: &str,
922        update: UpdateSkillScript,
923    ) -> anyhow::Result<SkillScriptRecord>;
924    async fn delete_script(&self, script_id: &str) -> anyhow::Result<()>;
925
926    // Discovery
927    async fn list_public_skills(&self) -> anyhow::Result<Vec<SkillRecord>>;
928    async fn star_skill(&self, skill_id: &str) -> anyhow::Result<()>;
929    async fn unstar_skill(&self, skill_id: &str) -> anyhow::Result<()>;
930    async fn list_starred_skills(&self) -> anyhow::Result<Vec<SkillRecord>>;
931    async fn clone_skill(&self, skill_id: &str) -> anyhow::Result<SkillRecord>;
932}
933
934// ========== Workflow Store ==========
935
936/// API response wrapper for workflow list endpoints.
937#[derive(Debug, Clone, Serialize, Deserialize)]
938pub struct WorkflowsListResponse {
939    pub workflows: Vec<WorkflowListItem>,
940    pub total: i64,
941}
942
943/// Lightweight workflow record for list endpoints — no definition payload.
944#[derive(Debug, Clone, Serialize, Deserialize)]
945pub struct WorkflowListItem {
946    pub id: String,
947    pub name: String,
948    #[serde(default)]
949    pub description: Option<String>,
950    #[serde(default)]
951    pub tags: Vec<String>,
952    #[serde(default)]
953    pub is_public: bool,
954    #[serde(default)]
955    pub is_template: bool,
956    #[serde(default)]
957    pub is_owner: bool,
958    #[serde(default)]
959    pub star_count: i32,
960    #[serde(default)]
961    pub clone_count: i32,
962    #[serde(default)]
963    pub is_starred: bool,
964    pub status: String,
965    pub step_count: usize,
966    pub created_at: chrono::DateTime<chrono::Utc>,
967    pub updated_at: chrono::DateTime<chrono::Utc>,
968}
969
970/// Full workflow record with definition payload.
971/// `definition` is `serde_json::Value` to avoid crate dependency on distri-workflow.
972/// Deserialize to `WorkflowDefinition` on demand.
973#[derive(Debug, Clone, Serialize, Deserialize)]
974pub struct WorkflowRecord {
975    pub id: String,
976    pub name: String,
977    pub description: Option<String>,
978    #[serde(default)]
979    /// The workflow definition as JSON. Deserialize to distri_workflow::WorkflowDefinition.
980    pub definition: serde_json::Value,
981    pub tags: Vec<String>,
982    pub is_public: bool,
983    pub is_template: bool,
984    pub star_count: i32,
985    pub clone_count: i32,
986    pub created_at: chrono::DateTime<chrono::Utc>,
987    pub updated_at: chrono::DateTime<chrono::Utc>,
988}
989
990/// Request to create a new workflow.
991#[derive(Debug, Clone, Serialize, Deserialize)]
992pub struct NewWorkflow {
993    pub name: String,
994    pub description: Option<String>,
995    #[serde(default)]
996    pub definition: serde_json::Value,
997    #[serde(default)]
998    pub tags: Vec<String>,
999    #[serde(default)]
1000    pub is_public: bool,
1001    #[serde(default)]
1002    pub is_template: bool,
1003}
1004
1005/// Partial update for a workflow.
1006#[derive(Debug, Clone, Serialize, Deserialize)]
1007pub struct UpdateWorkflow {
1008    pub name: Option<String>,
1009    pub description: Option<String>,
1010    pub definition: Option<serde_json::Value>,
1011    pub tags: Option<Vec<String>>,
1012    pub is_public: Option<bool>,
1013}
1014
1015/// Filter for listing workflows.
1016#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1017pub struct WorkflowFilter {
1018    pub is_public: Option<bool>,
1019    pub is_template: Option<bool>,
1020    pub tags: Option<Vec<String>>,
1021    pub search: Option<String>,
1022    pub limit: Option<i64>,
1023    pub offset: Option<i64>,
1024}
1025
1026#[async_trait]
1027pub trait WorkflowStore: Send + Sync {
1028    async fn list_workflows(&self, filter: WorkflowFilter) -> anyhow::Result<Vec<WorkflowRecord>>;
1029    async fn get_workflow(&self, id: &str) -> anyhow::Result<Option<WorkflowRecord>>;
1030    async fn create_workflow(&self, workflow: NewWorkflow) -> anyhow::Result<WorkflowRecord>;
1031    async fn update_workflow(
1032        &self,
1033        id: &str,
1034        update: UpdateWorkflow,
1035    ) -> anyhow::Result<WorkflowRecord>;
1036    async fn delete_workflow(&self, id: &str) -> anyhow::Result<()>;
1037
1038    // Discovery
1039    async fn list_public_workflows(&self) -> anyhow::Result<Vec<WorkflowRecord>>;
1040    async fn star_workflow(&self, workflow_id: &str) -> anyhow::Result<()>;
1041    async fn unstar_workflow(&self, workflow_id: &str) -> anyhow::Result<()>;
1042    async fn list_starred_workflows(&self) -> anyhow::Result<Vec<WorkflowRecord>>;
1043    async fn clone_workflow(&self, workflow_id: &str) -> anyhow::Result<WorkflowRecord>;
1044}
1045
1046// ─── Usage Service ──────────────────────────────────────────────────────────
1047
1048/// Current usage snapshot for a workspace/user.
1049#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1050pub struct UsageSnapshot {
1051    pub day_tokens: i64,
1052    pub week_tokens: i64,
1053    pub month_tokens: i64,
1054}
1055
1056/// Configured token limits for a workspace.
1057#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1058pub struct UsageLimits {
1059    pub daily_tokens: Option<i64>,
1060    pub weekly_tokens: Option<i64>,
1061    pub monthly_tokens: Option<i64>,
1062}
1063
1064/// Result of a rate limit check.
1065#[derive(Debug, Clone)]
1066pub enum UsageCheckResult {
1067    Allowed,
1068    Denied { reason: String },
1069}
1070
1071/// Trait for usage tracking, rate limiting, and workspace limit management.
1072///
1073/// OSS: can use a no-op or in-memory implementation.
1074/// Cloud: backed by Redis + Postgres with caching.
1075#[async_trait]
1076pub trait UsageService: Send + Sync {
1077    /// Check whether a request should be allowed based on all rate limits.
1078    /// Called by middleware before processing a request.
1079    /// `is_llm` indicates whether this is an LLM-consuming endpoint.
1080    /// `auth_source` is "jwt" or "api_key" for per-source analytics.
1081    async fn check_request(
1082        &self,
1083        workspace_id: &str,
1084        user_id: &str,
1085        is_llm: bool,
1086        auth_source: &str,
1087    ) -> UsageCheckResult;
1088
1089    /// Record token usage after a completed agent run.
1090    async fn record_usage(
1091        &self,
1092        workspace_id: &str,
1093        user_id: &str,
1094        tokens_used: i64,
1095    ) -> anyhow::Result<()>;
1096
1097    /// Get current usage snapshot for display.
1098    async fn get_usage(&self, workspace_id: &str, user_id: &str) -> anyhow::Result<UsageSnapshot>;
1099
1100    /// Get the configured limits for a workspace.
1101    async fn get_limits(&self, workspace_id: &str) -> anyhow::Result<UsageLimits>;
1102}
1103
1104/// No-op usage service for OSS / development.
1105/// Always allows requests, never records anything.
1106#[derive(Debug, Clone)]
1107pub struct NoOpUsageService;
1108
1109#[async_trait]
1110impl UsageService for NoOpUsageService {
1111    async fn check_request(
1112        &self,
1113        _workspace_id: &str,
1114        _user_id: &str,
1115        _is_llm: bool,
1116        _auth_source: &str,
1117    ) -> UsageCheckResult {
1118        UsageCheckResult::Allowed
1119    }
1120
1121    async fn record_usage(
1122        &self,
1123        _workspace_id: &str,
1124        _user_id: &str,
1125        _tokens_used: i64,
1126    ) -> anyhow::Result<()> {
1127        Ok(())
1128    }
1129
1130    async fn get_usage(
1131        &self,
1132        _workspace_id: &str,
1133        _user_id: &str,
1134    ) -> anyhow::Result<UsageSnapshot> {
1135        Ok(UsageSnapshot::default())
1136    }
1137
1138    async fn get_limits(&self, _workspace_id: &str) -> anyhow::Result<UsageLimits> {
1139        Ok(UsageLimits::default())
1140    }
1141}
1142
1143#[cfg(test)]
1144mod tests {
1145    use super::*;
1146
1147    #[test]
1148    fn test_skills_list_response_deserialize_cloud_format() {
1149        let json = r#"{"skills":[{"id":"abc","workspace_slug":"ws","name":"test","full_name":"ws/test","description":"desc","tags":["t"],"is_public":true,"is_system":false,"is_owner":true,"star_count":0,"clone_count":0,"is_starred":false,"created_at":"2026-01-01T00:00:00Z","updated_at":"2026-01-01T00:00:00Z"}]}"#;
1150        let resp: SkillsListResponse = serde_json::from_str(json).unwrap();
1151        assert_eq!(resp.skills.len(), 1);
1152        assert_eq!(resp.skills[0].name, "test");
1153        assert_eq!(resp.skills[0].workspace_slug, "ws");
1154        assert_eq!(resp.skills[0].full_name, "ws/test");
1155        assert!(resp.skills[0].is_public);
1156    }
1157
1158    #[test]
1159    fn test_skills_list_response_deserialize_defaults() {
1160        let json = r#"{"skills":[{"id":"abc","name":"test","created_at":"2026-01-01T00:00:00Z","updated_at":"2026-01-01T00:00:00Z"}]}"#;
1161        let resp: SkillsListResponse = serde_json::from_str(json).unwrap();
1162        assert_eq!(resp.skills[0].workspace_slug, "");
1163        assert_eq!(resp.skills[0].full_name, "");
1164        assert!(!resp.skills[0].is_public);
1165        assert!(!resp.skills[0].is_owner);
1166    }
1167
1168    #[test]
1169    fn test_skills_list_response_roundtrip() {
1170        let resp = SkillsListResponse {
1171            skills: vec![SkillListItem {
1172                id: "id1".into(),
1173                workspace_slug: "local".into(),
1174                name: "my_skill".into(),
1175                full_name: "local/my_skill".into(),
1176                description: Some("A skill".into()),
1177                tags: vec!["tag1".into()],
1178                is_public: false,
1179                is_system: false,
1180                is_owner: true,
1181                star_count: 5,
1182                clone_count: 2,
1183                is_starred: true,
1184                created_at: chrono::Utc::now(),
1185                updated_at: chrono::Utc::now(),
1186            }],
1187        };
1188        let json = serde_json::to_string(&resp).unwrap();
1189        let decoded: SkillsListResponse = serde_json::from_str(&json).unwrap();
1190        assert_eq!(decoded.skills[0].name, "my_skill");
1191        assert_eq!(decoded.skills[0].star_count, 5);
1192    }
1193}
distri_types/stores.rs

distri_types/
stores.rs
