Skip to main content

surfpool_types/
scenarios.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4use solana_clock::Slot;
5use uuid::Uuid;
6
7use crate::Idl;
8
9// ========================================
10// Core Scenarios Types
11// ========================================
12
13/// Defines how an account address should be determined
14#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
15#[serde(rename_all = "camelCase")]
16#[doc = "Defines how an account address should be determined"]
17pub enum AccountAddress {
18    /// A specific public key
19    #[doc = "A specific public key"]
20    Pubkey(String),
21    /// A Program Derived Address with seeds
22    #[doc = "A Program Derived Address with seeds"]
23    Pda {
24        program_id: String,
25        seeds: Vec<PdaSeed>,
26    },
27}
28
29/// Seeds used for PDA derivation
30#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
31#[serde(rename_all = "camelCase")]
32#[doc = "Seeds used for PDA derivation"]
33pub enum PdaSeed {
34    Pubkey(String),
35    String(String),
36    Bytes(Vec<u8>),
37    /// Reference to a property value
38    PropertyRef(String),
39}
40
41/// A reusable template for creating account overrides
42/// Values are mapped directly to IDL fields using dot notation (e.g., "agg.price", "expo")
43#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
44#[serde(rename_all = "camelCase")]
45pub struct OverrideTemplate {
46    /// Unique identifier for the template
47    pub id: String,
48    /// Human-readable name
49    pub name: String,
50    /// Description of what this template does
51    pub description: String,
52    /// Protocol this template is for (e.g., "Pyth", "Switchboard")
53    pub protocol: String,
54    /// IDL for the account structure - defines all available fields and types
55    pub idl: Idl,
56    /// How to determine the account address
57    pub address: AccountAddress,
58    /// Account type name from the IDL (e.g., "PriceAccount")
59    /// This specifies which account struct in the IDL to use
60    pub account_type: String,
61    pub properties: Vec<String>,
62    /// Tags for categorization and search
63    pub tags: Vec<String>,
64}
65
66impl OverrideTemplate {
67    pub fn new(
68        id: String,
69        name: String,
70        description: String,
71        protocol: String,
72        idl: Idl,
73        address: AccountAddress,
74        properties: Vec<String>,
75        account_type: String,
76    ) -> Self {
77        Self {
78            id,
79            name,
80            description,
81            protocol,
82            idl,
83            address,
84            account_type,
85            properties,
86            tags: Vec::new(),
87        }
88    }
89}
90
91/// A concrete instance of an override template with specific values
92#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
93#[serde(rename_all = "camelCase")]
94pub struct OverrideInstance {
95    /// Unique identifier for this instance
96    #[doc = "Unique identifier for the scenario"]
97    pub id: String,
98    /// Reference to the template being used
99    #[doc = "Reference to the template being used"]
100    pub template_id: String,
101    /// Values for the template properties (flat key-value map with dot notation, e.g., "price_message.price_value")
102    #[doc = "Values for the template properties (flat key-value map with dot notation, e.g., 'price_message.price_value')"]
103    pub values: HashMap<String, serde_json::Value>,
104    /// Relative slot when this override should be applied (relative to scenario registration slot)
105    #[doc = "Relative slot when this override should be applied (relative to scenario registration slot)"]
106    pub scenario_relative_slot: Slot,
107    /// Optional label for this instance
108    #[doc = "Optional label for this instance"]
109    pub label: Option<String>,
110    /// Whether this override is enabled
111    #[doc = "Whether this override is enabled"]
112    pub enabled: bool,
113    /// Whether to fetch fresh account data just before transaction execution
114    #[doc = "Whether to fetch fresh account data just before transaction execution"]
115    #[serde(default)]
116    pub fetch_before_use: bool,
117    /// Account address to override
118    #[doc = "Account address to override"]
119    pub account: AccountAddress,
120}
121
122impl OverrideInstance {
123    pub fn new(template_id: String, scenario_relative_slot: Slot, account: AccountAddress) -> Self {
124        Self {
125            id: Uuid::new_v4().to_string(),
126            template_id,
127            values: HashMap::new(),
128            scenario_relative_slot,
129            label: None,
130            enabled: true,
131            fetch_before_use: false,
132            account,
133        }
134    }
135
136    pub fn with_values(mut self, values: HashMap<String, serde_json::Value>) -> Self {
137        self.values = values;
138        self
139    }
140
141    pub fn with_label(mut self, label: String) -> Self {
142        self.label = Some(label);
143        self
144    }
145}
146
147/// A scenario containing a timeline of overrides
148#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
149#[serde(rename_all = "camelCase")]
150pub struct Scenario {
151    /// Unique identifier for the scenario
152    #[doc = "Unique identifier for the scenario"]
153    pub id: String,
154    /// Human-readable name
155    #[doc = "Human-readable name"]
156    pub name: String,
157    /// Description of this scenario
158    #[doc = "Description of this scenario"]
159    pub description: String,
160    /// List of override instances in this scenario
161    #[doc = "List of override instances in this scenario"]
162    pub overrides: Vec<OverrideInstance>,
163    /// Tags for categorization
164    #[doc = "Tags for categorization"]
165    pub tags: Vec<String>,
166}
167
168impl Scenario {
169    pub fn new(name: String, description: String) -> Self {
170        Self {
171            id: Uuid::new_v4().to_string(),
172            name,
173            description,
174            overrides: Vec::new(),
175            tags: Vec::new(),
176        }
177    }
178
179    pub fn add_override(&mut self, override_instance: OverrideInstance) {
180        self.overrides.push(override_instance);
181        // Sort by slot for efficient lookup
182        self.overrides.sort_by_key(|o| o.scenario_relative_slot);
183    }
184
185    pub fn remove_override(&mut self, override_id: &str) {
186        self.overrides.retain(|o| o.id != override_id);
187    }
188
189    pub fn get_overrides_for_slot(&self, slot: Slot) -> Vec<&OverrideInstance> {
190        self.overrides
191            .iter()
192            .filter(|o| o.enabled && o.scenario_relative_slot == slot)
193            .collect()
194    }
195}
196
197/// Configuration for scenario execution
198#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
199#[serde(rename_all = "camelCase")]
200pub struct ScenarioConfig {
201    /// Whether scenarios are enabled
202    pub enabled: bool,
203    /// Currently active scenario
204    pub active_scenario: Option<String>,
205    /// Whether to auto-save scenario changes
206    pub auto_save: bool,
207}
208
209impl Default for ScenarioConfig {
210    fn default() -> Self {
211        Self {
212            enabled: false,
213            active_scenario: None,
214            auto_save: true,
215        }
216    }
217}
218
219// ========================================
220// YAML Template File Types
221// ========================================
222
223/// YAML representation of an override template loaded from file
224/// References an external IDL file via idl_file_path
225#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
226pub struct YamlOverrideTemplateFile {
227    pub id: String,
228    pub name: String,
229    pub description: String,
230    pub protocol: String,
231    pub version: String,
232    pub account_type: String,
233    pub properties: Vec<String>,
234    pub idl_file_path: String,
235    pub address: YamlAccountAddress,
236    #[serde(default)]
237    pub tags: Vec<String>,
238}
239
240impl YamlOverrideTemplateFile {
241    /// Convert file-based template to runtime OverrideTemplate with loaded IDL
242    pub fn to_override_template(self, idl: Idl) -> OverrideTemplate {
243        OverrideTemplate {
244            id: self.id,
245            name: self.name,
246            description: self.description,
247            protocol: self.protocol,
248            idl,
249            address: self.address.into(),
250            account_type: self.account_type,
251            properties: self.properties,
252            tags: self.tags,
253        }
254    }
255}
256
257/// Collection of override templates sharing the same IDL
258/// Used when one YAML file defines multiple templates (e.g., multiple Pyth price feeds)
259#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
260pub struct YamlOverrideTemplateCollection {
261    /// Protocol these templates are for
262    pub protocol: String,
263    /// Version identifier
264    pub version: String,
265    /// Path to shared IDL file
266    pub idl_file_path: String,
267    /// Common tags for all templates
268    #[serde(default)]
269    pub tags: Vec<String>,
270    /// The templates
271    pub templates: Vec<YamlOverrideTemplateEntry>,
272}
273
274/// Individual template entry in a collection
275#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
276pub struct YamlOverrideTemplateEntry {
277    pub id: String,
278    pub name: String,
279    pub description: String,
280    pub idl_account_name: String,
281    pub properties: Vec<String>,
282    pub address: YamlAccountAddress,
283}
284
285impl YamlOverrideTemplateCollection {
286    /// Convert collection to runtime OverrideTemplates with loaded IDL
287    pub fn to_override_templates(self, idl: Idl) -> Vec<OverrideTemplate> {
288        self.templates
289            .into_iter()
290            .map(|entry| OverrideTemplate {
291                id: entry.id,
292                name: entry.name,
293                description: entry.description,
294                protocol: self.protocol.clone(),
295                idl: idl.clone(),
296                address: entry.address.into(),
297                account_type: entry.idl_account_name,
298                properties: entry.properties,
299                tags: self.tags.clone(),
300            })
301            .collect()
302    }
303}
304
305/// YAML representation of an override template with embedded IDL
306/// Used for RPC methods where file access is not available
307#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
308pub struct YamlOverrideTemplate {
309    pub id: String,
310    pub name: String,
311    pub description: String,
312    pub protocol: String,
313    pub version: String,
314    pub account_type: String,
315    pub idl: Idl,
316    pub address: YamlAccountAddress,
317    pub properties: Vec<String>,
318    #[serde(default)]
319    pub tags: Vec<String>,
320}
321
322impl YamlOverrideTemplate {
323    /// Convert to runtime OverrideTemplate
324    pub fn to_override_template(self) -> OverrideTemplate {
325        OverrideTemplate {
326            id: self.id,
327            name: self.name,
328            description: self.description,
329            protocol: self.protocol,
330            idl: self.idl,
331            address: self.address.into(),
332            account_type: self.account_type,
333            properties: self.properties,
334            tags: self.tags,
335        }
336    }
337}
338
339/// YAML representation of account address
340#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
341#[serde(tag = "type", rename_all = "lowercase")]
342pub enum YamlAccountAddress {
343    Pubkey {
344        #[serde(default)]
345        value: Option<String>,
346    },
347    Pda {
348        program_id: String,
349        seeds: Vec<YamlPdaSeed>,
350    },
351}
352
353impl From<YamlAccountAddress> for AccountAddress {
354    fn from(yaml: YamlAccountAddress) -> Self {
355        match yaml {
356            YamlAccountAddress::Pubkey { value } => {
357                AccountAddress::Pubkey(value.unwrap_or_default())
358            }
359            YamlAccountAddress::Pda { program_id, seeds } => AccountAddress::Pda {
360                program_id,
361                seeds: seeds.into_iter().map(|s| s.into()).collect(),
362            },
363        }
364    }
365}
366
367/// YAML representation of PDA seeds
368#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
369#[serde(tag = "type", rename_all = "snake_case")]
370pub enum YamlPdaSeed {
371    String { value: String },
372    Bytes { value: Vec<u8> },
373    Pubkey { value: String },
374    PropertyRef { value: String },
375}
376
377impl From<YamlPdaSeed> for PdaSeed {
378    fn from(yaml: YamlPdaSeed) -> Self {
379        match yaml {
380            YamlPdaSeed::String { value } => PdaSeed::String(value),
381            YamlPdaSeed::Bytes { value } => PdaSeed::Bytes(value),
382            YamlPdaSeed::Pubkey { value } => PdaSeed::Pubkey(value),
383            YamlPdaSeed::PropertyRef { value } => PdaSeed::PropertyRef(value),
384        }
385    }
386}