Skip to main content

surfpool_types/
scenarios.rs

1use std::{collections::HashMap, str::FromStr};
2
3use serde::{Deserialize, Serialize};
4use solana_clock::Slot;
5use solana_pubkey::Pubkey;
6use uuid::Uuid;
7
8use crate::Idl;
9
10// ========================================
11// Constants Types (for UI comboboxes and LLM choices)
12// ========================================
13
14/// A single option within a constant definition
15/// Used to define selectable values like AMM configs or well-known tokens
16#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
17#[serde(rename_all = "camelCase")]
18pub struct ConstantOption {
19    /// Unique identifier for this option (e.g., "standard", "sol", "usdc")
20    pub id: String,
21    /// Human-readable label shown in UI (e.g., "Standard (25 bps)", "SOL (Wrapped)")
22    pub label: String,
23    /// Description explaining when to use this option
24    #[serde(default)]
25    pub description: Option<String>,
26    /// The actual value (typically a pubkey string)
27    pub value: String,
28    /// Additional metadata for context (e.g., decimals, tick_spacing, fee rates)
29    #[serde(default)]
30    pub metadata: HashMap<String, serde_json::Value>,
31}
32
33/// A constant definition containing multiple selectable options
34/// Used to define things like AMM fee tiers, well-known tokens, etc.
35#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
36#[serde(rename_all = "camelCase")]
37pub struct ConstantDefinition {
38    /// Human-readable label for this constant type (e.g., "Fee Tier", "Token")
39    pub label: String,
40    /// Description of what this constant represents
41    #[serde(default)]
42    pub description: Option<String>,
43    /// The available options to choose from
44    pub options: Vec<ConstantOption>,
45}
46
47impl ConstantDefinition {
48    /// Get an option by its ID
49    pub fn get_option(&self, id: &str) -> Option<&ConstantOption> {
50        self.options.iter().find(|o| o.id == id)
51    }
52
53    /// Get the value for an option by ID
54    pub fn get_value(&self, id: &str) -> Option<&str> {
55        self.get_option(id).map(|o| o.value.as_str())
56    }
57}
58
59// ========================================
60// Core Scenarios Types
61// ========================================
62
63/// Defines how an account address should be determined
64#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
65#[serde(rename_all = "camelCase")]
66#[doc = "Defines how an account address should be determined"]
67pub enum AccountAddress {
68    /// A specific public key
69    #[doc = "A specific public key"]
70    Pubkey(String),
71    /// A Program Derived Address with seeds
72    #[doc = "A Program Derived Address with seeds"]
73    #[serde(rename_all = "camelCase")]
74    Pda {
75        program_id: String,
76        seeds: Vec<PdaSeed>,
77    },
78}
79
80/// Seeds used for PDA derivation
81#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
82#[serde(rename_all = "camelCase")]
83#[doc = "Seeds used for PDA derivation"]
84pub enum PdaSeed {
85    Pubkey(String),
86    String(String),
87    Bytes(Vec<u8>),
88    /// Reference to a property value
89    PropertyRef(String),
90    /// A u16 value converted to big-endian bytes (useful for config indices)
91    U16Be(u16),
92    /// Reference to a property that should be converted to u16 big-endian bytes
93    U16BeRef(String),
94    /// A u16 value converted to little-endian bytes (useful for Pyth shard IDs)
95    U16Le(u16),
96    /// Reference to a property that's a 32-byte hex string (e.g., Pyth feed ID)
97    Bytes32Ref(String),
98    /// A nested PDA derivation - derives a PDA from inner seeds and uses it as the seed
99    #[serde(rename_all = "camelCase")]
100    DerivedPda {
101        program_id: String,
102        seeds: Vec<PdaSeed>,
103    },
104}
105
106impl PdaSeed {
107    /// Convert a seed to bytes, optionally using values for PropertyRef resolution
108    pub fn to_bytes(&self, values: Option<&HashMap<String, serde_json::Value>>) -> Option<Vec<u8>> {
109        match self {
110            PdaSeed::Pubkey(pk_str) => Pubkey::from_str(pk_str)
111                .ok()
112                .map(|pk| pk.to_bytes().to_vec()),
113            PdaSeed::String(s) => Some(s.as_bytes().to_vec()),
114            PdaSeed::Bytes(b) => Some(b.clone()),
115            PdaSeed::PropertyRef(prop) => {
116                values?.get(prop).and_then(|v| {
117                    // Handle string values (could be pubkey or raw string)
118                    if let Some(s) = v.as_str() {
119                        if let Ok(pk) = Pubkey::from_str(s) {
120                            return Some(pk.to_bytes().to_vec());
121                        }
122                        return Some(s.as_bytes().to_vec());
123                    }
124                    // Handle numeric values (u64)
125                    if let Some(n) = v.as_u64() {
126                        return Some(n.to_le_bytes().to_vec());
127                    }
128                    None
129                })
130            }
131            PdaSeed::U16Be(n) => Some(n.to_be_bytes().to_vec()),
132            PdaSeed::U16BeRef(prop) => {
133                values?.get(prop).and_then(|v| {
134                    // Handle numeric values - convert to u16 big-endian
135                    if let Some(n) = v.as_u64() {
136                        let n16 = u16::try_from(n).ok()?;
137                        return Some(n16.to_be_bytes().to_vec());
138                    }
139                    None
140                })
141            }
142            PdaSeed::U16Le(n) => Some(n.to_le_bytes().to_vec()),
143            PdaSeed::Bytes32Ref(prop) => {
144                values?.get(prop).and_then(|v| {
145                    // Handle hex string values (e.g., "0xef0d8b6f..." for Pyth feed IDs)
146                    if let Some(s) = v.as_str() {
147                        // Remove 0x prefix if present
148                        let hex_str = s.strip_prefix("0x").unwrap_or(s);
149                        // Parse as 32-byte hex
150                        if let Ok(bytes) = hex::decode(hex_str) {
151                            if bytes.len() == 32 {
152                                return Some(bytes);
153                            }
154                        }
155                    }
156                    None
157                })
158            }
159            PdaSeed::DerivedPda { program_id, seeds } => {
160                // Derive a nested PDA and use its pubkey as the seed
161                let program_pubkey = Pubkey::from_str(program_id).ok()?;
162
163                // Convert inner seeds to bytes
164                let seed_bytes: Vec<Vec<u8>> = seeds
165                    .iter()
166                    .filter_map(|seed| seed.to_bytes(values))
167                    .collect();
168
169                // Ensure all seeds were converted successfully
170                if seed_bytes.len() != seeds.len() {
171                    return None;
172                }
173
174                // Create seed slices for find_program_address
175                let seed_slices: Vec<&[u8]> = seed_bytes.iter().map(|s| s.as_slice()).collect();
176
177                // Derive the nested PDA
178                let (pda, _bump) = Pubkey::find_program_address(&seed_slices, &program_pubkey);
179                Some(pda.to_bytes().to_vec())
180            }
181        }
182    }
183}
184
185impl AccountAddress {
186    /// Resolve the account address to a Pubkey
187    /// For PDA addresses, this derives the address from the program_id and seeds
188    /// For PropertyRef seeds, values map is used to resolve the reference
189    pub fn resolve(&self, values: Option<&HashMap<String, serde_json::Value>>) -> Option<Pubkey> {
190        match self {
191            AccountAddress::Pubkey(pubkey_str) => Pubkey::from_str(pubkey_str).ok(),
192            AccountAddress::Pda { program_id, seeds } => {
193                let program_pubkey = Pubkey::from_str(program_id).ok()?;
194
195                // Convert all seeds to bytes
196                let seed_bytes: Vec<Vec<u8>> = seeds
197                    .iter()
198                    .filter_map(|seed| seed.to_bytes(values))
199                    .collect();
200
201                // Ensure all seeds were converted successfully
202                if seed_bytes.len() != seeds.len() {
203                    return None;
204                }
205
206                // Create seed slices for find_program_address
207                let seed_slices: Vec<&[u8]> = seed_bytes.iter().map(|s| s.as_slice()).collect();
208
209                // Derive the PDA
210                let (pda, _bump) = Pubkey::find_program_address(&seed_slices, &program_pubkey);
211                Some(pda)
212            }
213        }
214    }
215
216    /// Resolve the account address to a Pubkey without any values for PropertyRef
217    /// This is a convenience method when no PropertyRef seeds are expected
218    pub fn resolve_simple(&self) -> Option<Pubkey> {
219        self.resolve(None)
220    }
221
222    /// Get the names of values that are referenced by PDA seeds
223    /// These should be filtered out when applying account data overrides
224    /// since they're only used for address derivation, not account data modification
225    pub fn get_pda_seed_references(&self) -> Vec<String> {
226        match self {
227            AccountAddress::Pubkey(_) => vec![],
228            AccountAddress::Pda { seeds, .. } => {
229                let mut refs = Vec::new();
230                Self::collect_seed_references(seeds, &mut refs);
231                refs
232            }
233        }
234    }
235
236    /// Recursively collect property references from seeds
237    fn collect_seed_references(seeds: &[PdaSeed], refs: &mut Vec<String>) {
238        for seed in seeds {
239            match seed {
240                PdaSeed::PropertyRef(name) => refs.push(name.clone()),
241                PdaSeed::U16BeRef(name) => refs.push(name.clone()),
242                PdaSeed::Bytes32Ref(name) => refs.push(name.clone()),
243                PdaSeed::DerivedPda { seeds: inner, .. } => {
244                    Self::collect_seed_references(inner, refs);
245                }
246                _ => {}
247            }
248        }
249    }
250}
251
252/// The type of a property - determines how it's rendered in the UI
253#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
254#[serde(rename_all = "snake_case")]
255pub enum PropertyKind {
256    /// A regular field from the IDL (default behavior)
257    #[default]
258    Field,
259    /// A reference to a constant definition (renders as dropdown/combobox in UI)
260    ConstantRef,
261}
262
263/// Defines a property in a template with full metadata
264#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
265#[serde(rename_all = "snake_case")]
266pub struct Property {
267    /// The path to the field in the IDL (e.g., "liquidity", "fees.swap_fee_numerator")
268    pub path: String,
269    /// The type of property - determines rendering behavior
270    #[serde(default, rename = "type")]
271    pub kind: PropertyKind,
272    /// Human-readable label for the UI (optional, defaults to path)
273    #[serde(default, skip_serializing_if = "Option::is_none")]
274    pub label: Option<String>,
275    /// Description of the field (optional, can come from IDL)
276    #[serde(default, skip_serializing_if = "Option::is_none")]
277    pub description: Option<String>,
278    /// For constant_ref type: the name of the constant definition to use
279    #[serde(default, skip_serializing_if = "Option::is_none")]
280    pub constant: Option<String>,
281}
282
283impl Property {
284    /// Create a new field property
285    pub fn field(path: impl Into<String>) -> Self {
286        Self {
287            path: path.into(),
288            kind: PropertyKind::Field,
289            label: None,
290            description: None,
291            constant: None,
292        }
293    }
294
295    /// Create a new constant_ref property
296    pub fn constant_ref(path: impl Into<String>, constant: impl Into<String>) -> Self {
297        Self {
298            path: path.into(),
299            kind: PropertyKind::ConstantRef,
300            label: None,
301            description: None,
302            constant: Some(constant.into()),
303        }
304    }
305
306    /// Check if this is a constant reference
307    pub fn is_constant_ref(&self) -> bool {
308        matches!(self.kind, PropertyKind::ConstantRef)
309    }
310
311    /// Get the constant name if this is a constant reference
312    pub fn constant_name(&self) -> Option<&str> {
313        self.constant.as_deref()
314    }
315
316    /// Get the display label (falls back to path if no label set)
317    pub fn display_label(&self) -> &str {
318        self.label.as_deref().unwrap_or(&self.path)
319    }
320}
321
322/// Legacy tagged property representation.
323///
324/// Use [`Property`] for new code. This type remains deserializable for callers
325/// that still send the old tagged JSON shape, such as
326/// `{"type":"field","name":"price"}`.
327#[deprecated(
328    since = "1.5.0",
329    note = "use Property; PropertyType remains only for legacy tagged JSON compatibility"
330)]
331#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
332#[serde(rename_all = "snake_case", tag = "type")]
333pub enum PropertyType {
334    /// A regular field from the IDL (default behavior)
335    Field { name: String },
336    /// A reference to a constant definition (renders as combobox in UI)
337    ConstantRef { name: String, constant: String },
338}
339
340#[allow(deprecated)]
341impl PropertyType {
342    /// Get the property name regardless of type
343    pub fn name(&self) -> &str {
344        match self {
345            PropertyType::Field { name } => name,
346            PropertyType::ConstantRef { name, .. } => name,
347        }
348    }
349
350    /// Check if this is a constant reference
351    pub fn is_constant_ref(&self) -> bool {
352        matches!(self, PropertyType::ConstantRef { .. })
353    }
354
355    /// Get the constant name if this is a constant reference
356    pub fn constant_name(&self) -> Option<&str> {
357        match self {
358            PropertyType::ConstantRef { constant, .. } => Some(constant),
359            _ => None,
360        }
361    }
362}
363
364#[allow(deprecated)]
365impl From<PropertyType> for Property {
366    fn from(pt: PropertyType) -> Self {
367        match pt {
368            PropertyType::Field { name } => Property::field(name),
369            PropertyType::ConstantRef { name, constant } => Property::constant_ref(name, constant),
370        }
371    }
372}
373
374/// A reusable template for creating account overrides
375/// Values are mapped directly to IDL fields using dot notation (e.g., "agg.price", "expo")
376#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
377#[serde(rename_all = "camelCase")]
378pub struct OverrideTemplate {
379    /// Unique identifier for the template
380    pub id: String,
381    /// Human-readable name
382    pub name: String,
383    /// Description of what this template does
384    pub description: String,
385    /// Protocol this template is for (e.g., "Pyth", "Switchboard")
386    pub protocol: String,
387    /// IDL for the account structure - defines all available fields and types
388    pub idl: Idl,
389    /// How to determine the account address
390    pub address: AccountAddress,
391    /// Account type name from the IDL (e.g., "PriceAccount")
392    /// This specifies which account struct in the IDL to use
393    pub account_type: String,
394    /// List of editable properties with full metadata
395    pub properties: Vec<Property>,
396    /// Protocol-specific constants (e.g., AMM configs, well-known tokens)
397    #[serde(default)]
398    pub constants: HashMap<String, ConstantDefinition>,
399    /// Tags for categorization and search
400    pub tags: Vec<String>,
401    /// Optional context/instructions specifically for LLMs using this template
402    /// This helps LLMs understand how to correctly use the template
403    #[serde(default, skip_serializing_if = "Option::is_none")]
404    pub llm_context: Option<String>,
405}
406
407impl OverrideTemplate {
408    pub fn new(
409        id: String,
410        name: String,
411        description: String,
412        protocol: String,
413        idl: Idl,
414        address: AccountAddress,
415        properties: Vec<Property>,
416        account_type: String,
417    ) -> Self {
418        Self {
419            id,
420            name,
421            description,
422            protocol,
423            idl,
424            address,
425            account_type,
426            properties,
427            constants: HashMap::new(),
428            tags: Vec::new(),
429            llm_context: None,
430        }
431    }
432
433    /// Get a constant definition by name
434    pub fn get_constant(&self, name: &str) -> Option<&ConstantDefinition> {
435        self.constants.get(name)
436    }
437
438    /// Check if a property is a constant reference
439    pub fn is_property_constant_ref(&self, property_name: &str) -> bool {
440        self.properties
441            .iter()
442            .any(|p| p.path == property_name && p.is_constant_ref())
443    }
444
445    /// Get the constant definition for a property if it's a constant reference
446    pub fn get_property_constant(&self, property_name: &str) -> Option<&ConstantDefinition> {
447        self.properties
448            .iter()
449            .find(|p| p.path == property_name)
450            .and_then(|p| p.constant_name())
451            .and_then(|const_name| self.constants.get(const_name))
452    }
453
454    /// Get a property by its path
455    pub fn get_property(&self, path: &str) -> Option<&Property> {
456        self.properties.iter().find(|p| p.path == path)
457    }
458
459    /// Get the list of property paths (for backward compatibility)
460    pub fn property_paths(&self) -> Vec<&str> {
461        self.properties.iter().map(|p| p.path.as_str()).collect()
462    }
463}
464
465/// A concrete instance of an override template with specific values
466#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
467#[serde(rename_all = "camelCase")]
468pub struct OverrideInstance {
469    /// Unique identifier for this instance (UUID v4)
470    #[schemars(description = "Unique identifier for this instance (UUID v4 format)")]
471    pub id: String,
472    /// Reference to the template being used - MUST match a template id from get_override_templates
473    #[schemars(
474        description = "Template ID from get_override_templates (e.g., 'raydium-clmm-custom', 'kamino-obligation-health')"
475    )]
476    pub template_id: String,
477    /// Values for the template properties as a JSON object (NOT a string)
478    #[schemars(
479        description = "JSON object mapping property names to values. Keys must be from template.properties. Example: {\"liquidity\": 1000000, \"sqrt_price_x64\": 18446744073709551616}"
480    )]
481    pub values: HashMap<String, serde_json::Value>,
482    /// Relative slot when this override should be applied (1 = 400ms after registration)
483    #[schemars(
484        description = "Slot offset from scenario registration (integer, e.g., 1, 2, 3). Each slot is ~400ms."
485    )]
486    pub scenario_relative_slot: Slot,
487    /// Optional human-readable label for this instance
488    #[schemars(description = "Human-readable label describing what this override does")]
489    pub label: Option<String>,
490    /// Whether this override is enabled
491    #[schemars(description = "Whether this override is active (true/false)")]
492    pub enabled: bool,
493    /// Whether to fetch fresh account data just before transaction execution
494    #[schemars(
495        description = "If true, fetches fresh account data from mainnet before applying override"
496    )]
497    #[serde(default)]
498    pub fetch_before_use: bool,
499    /// Account address to override - use pubkey for known addresses or pda for derived addresses
500    #[schemars(
501        description = "Account address: either {\"pubkey\": \"base58_address\"} or {\"pda\": {\"programId\": \"...\", \"seeds\": [...]}}"
502    )]
503    pub account: AccountAddress,
504}
505
506impl OverrideInstance {
507    pub fn new(template_id: String, scenario_relative_slot: Slot, account: AccountAddress) -> Self {
508        Self {
509            id: Uuid::new_v4().to_string(),
510            template_id,
511            values: HashMap::new(),
512            scenario_relative_slot,
513            label: None,
514            enabled: true,
515            fetch_before_use: false,
516            account,
517        }
518    }
519
520    pub fn with_values(mut self, values: HashMap<String, serde_json::Value>) -> Self {
521        self.values = values;
522        self
523    }
524
525    pub fn with_label(mut self, label: String) -> Self {
526        self.label = Some(label);
527        self
528    }
529}
530
531/// A scenario containing a timeline of overrides
532#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
533#[serde(rename_all = "camelCase")]
534pub struct Scenario {
535    /// Unique identifier for the scenario (UUID v4 format)
536    #[schemars(
537        description = "Unique identifier for the scenario (UUID v4 format, e.g., '550e8400-e29b-41d4-a716-446655440000')"
538    )]
539    pub id: String,
540    /// Human-readable name
541    #[schemars(description = "Human-readable name for the scenario")]
542    pub name: String,
543    /// Description of this scenario
544    #[schemars(description = "Description of what this scenario does")]
545    pub description: String,
546    /// List of override instances in this scenario - MUST be an array, NOT a string
547    #[schemars(
548        description = "Array of override instances. IMPORTANT: This must be a JSON array [], not a JSON string. Each element is an OverrideInstance object."
549    )]
550    pub overrides: Vec<OverrideInstance>,
551    /// Tags for categorization
552    #[schemars(
553        description = "Array of string tags for categorization (e.g., ['liquidation', 'arbitrage'])"
554    )]
555    pub tags: Vec<String>,
556}
557
558impl Scenario {
559    pub fn new(name: String, description: String) -> Self {
560        Self {
561            id: Uuid::new_v4().to_string(),
562            name,
563            description,
564            overrides: Vec::new(),
565            tags: Vec::new(),
566        }
567    }
568
569    pub fn add_override(&mut self, override_instance: OverrideInstance) {
570        self.overrides.push(override_instance);
571        // Sort by slot for efficient lookup
572        self.overrides.sort_by_key(|o| o.scenario_relative_slot);
573    }
574
575    pub fn remove_override(&mut self, override_id: &str) {
576        self.overrides.retain(|o| o.id != override_id);
577    }
578
579    pub fn get_overrides_for_slot(&self, slot: Slot) -> Vec<&OverrideInstance> {
580        self.overrides
581            .iter()
582            .filter(|o| o.enabled && o.scenario_relative_slot == slot)
583            .collect()
584    }
585}
586
587/// Configuration for scenario execution
588#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
589#[serde(rename_all = "camelCase")]
590pub struct ScenarioConfig {
591    /// Whether scenarios are enabled
592    pub enabled: bool,
593    /// Currently active scenario
594    pub active_scenario: Option<String>,
595    /// Whether to auto-save scenario changes
596    pub auto_save: bool,
597}
598
599impl Default for ScenarioConfig {
600    fn default() -> Self {
601        Self {
602            enabled: false,
603            active_scenario: None,
604            auto_save: true,
605        }
606    }
607}
608
609// ========================================
610// YAML Template File Types
611// ========================================
612
613/// YAML representation of an override template loaded from file
614/// References an external IDL file via idl_file_path
615#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
616pub struct YamlOverrideTemplateFile {
617    pub id: String,
618    pub name: String,
619    pub description: String,
620    pub protocol: String,
621    pub version: String,
622    pub account_type: String,
623    #[serde(default)]
624    pub properties: Vec<YamlProperty>,
625    #[serde(default)]
626    pub constants: HashMap<String, YamlConstantDefinition>,
627    pub idl_file_path: String,
628    pub address: YamlAccountAddress,
629    #[serde(default)]
630    pub tags: Vec<String>,
631    /// Optional context/instructions specifically for LLMs using this template
632    #[serde(default)]
633    pub llm_context: Option<String>,
634}
635
636impl YamlOverrideTemplateFile {
637    /// Convert file-based template to runtime OverrideTemplate with loaded IDL
638    pub fn to_override_template(self, idl: Idl) -> OverrideTemplate {
639        OverrideTemplate {
640            id: self.id,
641            name: self.name,
642            description: self.description,
643            protocol: self.protocol,
644            idl,
645            address: self.address.into(),
646            account_type: self.account_type,
647            properties: self.properties.into_iter().map(Into::into).collect(),
648            constants: self
649                .constants
650                .into_iter()
651                .map(|(k, v)| (k, v.into()))
652                .collect(),
653            tags: self.tags,
654            llm_context: self.llm_context,
655        }
656    }
657}
658
659/// YAML representation of a constant option
660#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
661pub struct YamlConstantOption {
662    /// Unique identifier for this option
663    pub id: String,
664    /// Human-readable label
665    pub label: String,
666    /// Description of when to use this option
667    #[serde(default)]
668    pub description: Option<String>,
669    /// The actual value (typically a pubkey)
670    pub value: String,
671    /// Additional metadata
672    #[serde(default)]
673    pub metadata: HashMap<String, serde_json::Value>,
674}
675
676impl From<YamlConstantOption> for ConstantOption {
677    fn from(yaml: YamlConstantOption) -> Self {
678        ConstantOption {
679            id: yaml.id,
680            label: yaml.label,
681            description: yaml.description,
682            value: yaml.value,
683            metadata: yaml.metadata,
684        }
685    }
686}
687
688/// Source for constant options - either inline or from verified tokens registry
689#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
690#[serde(untagged)]
691pub enum YamlConstantSource {
692    /// Inline options defined in the YAML
693    Inline { options: Vec<YamlConstantOption> },
694    /// Reference to the verified tokens registry
695    TokensRef {
696        /// Type of reference - currently only "verified_tokens" is supported
697        source: String,
698        /// Optional filter for which tokens to include (e.g., by tags like "major", "stable")
699        #[serde(default)]
700        filter_tags: Vec<String>,
701        /// Optional limit on number of tokens to include
702        #[serde(default)]
703        limit: Option<usize>,
704    },
705}
706
707/// YAML representation of a constant definition
708#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
709pub struct YamlConstantDefinition {
710    /// Human-readable label for this constant type
711    pub label: String,
712    /// Description of what this constant represents
713    #[serde(default)]
714    pub description: Option<String>,
715    /// The source of options - either inline or from a reference
716    #[serde(flatten)]
717    pub source: YamlConstantSource,
718}
719
720impl YamlConstantDefinition {
721    /// Convert to runtime ConstantDefinition, resolving verified tokens references
722    pub fn to_constant_definition(self) -> ConstantDefinition {
723        let options = match self.source {
724            YamlConstantSource::Inline { options } => options.into_iter().map(Into::into).collect(),
725            YamlConstantSource::TokensRef {
726                source,
727                filter_tags,
728                limit,
729            } => {
730                if source == "verified_tokens" {
731                    use crate::verified_tokens::VERIFIED_TOKENS_BY_SYMBOL;
732
733                    let mut tokens: Vec<_> = VERIFIED_TOKENS_BY_SYMBOL
734                        .iter()
735                        .filter(|(_, _token)| {
736                            // If no filter tags specified, include all tokens
737                            if filter_tags.is_empty() {
738                                return true;
739                            }
740                            // Check if token has any of the filter tags
741                            // The tags are stored in the CSV but not parsed into TokenInfo yet
742                            // For now, we'll include all tokens when filter is specified
743                            // TODO: Parse tags from CSV into TokenInfo struct
744                            true
745                        })
746                        .map(|(symbol, token)| ConstantOption {
747                            id: symbol.to_lowercase(),
748                            label: format!("{} ({})", token.symbol, token.name),
749                            description: Some(token.name.clone()),
750                            value: token.address.clone(),
751                            metadata: {
752                                let mut meta = HashMap::new();
753                                meta.insert(
754                                    "symbol".to_string(),
755                                    serde_json::Value::String(token.symbol.clone()),
756                                );
757                                meta.insert(
758                                    "decimals".to_string(),
759                                    serde_json::Value::Number(token.decimals.into()),
760                                );
761                                if let Some(ref logo) = token.logo_uri {
762                                    meta.insert(
763                                        "logo_uri".to_string(),
764                                        serde_json::Value::String(logo.clone()),
765                                    );
766                                }
767                                meta
768                            },
769                        })
770                        .collect();
771
772                    // Sort by symbol for consistent ordering
773                    tokens.sort_by(|a, b| a.id.cmp(&b.id));
774
775                    // Apply limit if specified
776                    if let Some(limit) = limit {
777                        tokens.truncate(limit);
778                    }
779
780                    tokens
781                } else {
782                    // Unknown source type - return empty options
783                    Vec::new()
784                }
785            }
786        };
787
788        ConstantDefinition {
789            label: self.label,
790            description: self.description,
791            options,
792        }
793    }
794}
795
796// Keep From impl for backward compatibility but use the new method
797impl From<YamlConstantDefinition> for ConstantDefinition {
798    fn from(yaml: YamlConstantDefinition) -> Self {
799        yaml.to_constant_definition()
800    }
801}
802
803/// YAML representation of a property (supports both simple string and full object format)
804#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
805#[serde(untagged)]
806pub enum YamlProperty {
807    /// Simple string format: just the path (e.g., "liquidity", "fees.swap_fee_numerator")
808    Simple(String),
809    /// Full object format with all metadata
810    Full {
811        /// The path to the field in the IDL
812        path: String,
813        /// The type of property: "field" (default) or "constant_ref"
814        #[serde(default, rename = "type")]
815        kind: Option<String>,
816        /// Human-readable label for the UI (optional)
817        #[serde(default)]
818        label: Option<String>,
819        /// Description of the field (optional)
820        #[serde(default)]
821        description: Option<String>,
822        /// For constant_ref type: the name of the constant definition to use
823        #[serde(default)]
824        constant: Option<String>,
825    },
826}
827
828impl From<YamlProperty> for Property {
829    fn from(yaml: YamlProperty) -> Self {
830        match yaml {
831            YamlProperty::Simple(path) => Property::field(path),
832            YamlProperty::Full {
833                path,
834                kind,
835                label,
836                description,
837                constant,
838            } => {
839                let kind = match kind.as_deref() {
840                    Some("constant_ref") => PropertyKind::ConstantRef,
841                    _ => PropertyKind::Field,
842                };
843                Property {
844                    path,
845                    kind,
846                    label,
847                    description,
848                    constant,
849                }
850            }
851        }
852    }
853}
854
855/// YAML representation of a typed property (deprecated, for backward compatibility)
856#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
857#[serde(tag = "type", rename_all = "snake_case")]
858pub enum YamlPropertyType {
859    /// A regular field from the IDL
860    Field { name: String },
861    /// A reference to a constant definition
862    ConstantRef { name: String, constant: String },
863}
864
865impl From<YamlPropertyType> for Property {
866    fn from(yaml: YamlPropertyType) -> Self {
867        match yaml {
868            YamlPropertyType::Field { name } => Property::field(name),
869            YamlPropertyType::ConstantRef { name, constant } => {
870                Property::constant_ref(name, constant)
871            }
872        }
873    }
874}
875
876#[allow(deprecated)]
877impl From<YamlPropertyType> for PropertyType {
878    fn from(yaml: YamlPropertyType) -> Self {
879        match yaml {
880            YamlPropertyType::Field { name } => PropertyType::Field { name },
881            YamlPropertyType::ConstantRef { name, constant } => {
882                PropertyType::ConstantRef { name, constant }
883            }
884        }
885    }
886}
887
888/// Collection of override templates sharing the same IDL
889/// Used when one YAML file defines multiple templates (e.g., multiple Pyth price feeds)
890#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
891pub struct YamlOverrideTemplateCollection {
892    /// Protocol these templates are for
893    pub protocol: String,
894    /// Version identifier
895    pub version: String,
896    /// Account type name from the IDL (optional, can be overridden per template)
897    #[serde(default)]
898    pub account_type: Option<String>,
899    /// Path to shared IDL file
900    pub idl_file_path: String,
901    /// Common tags for all templates
902    #[serde(default)]
903    pub tags: Vec<String>,
904    /// Protocol-specific constants shared by all templates in this collection
905    #[serde(default)]
906    pub constants: HashMap<String, YamlConstantDefinition>,
907    /// The templates
908    pub templates: Vec<YamlOverrideTemplateEntry>,
909}
910
911/// Individual template entry in a collection
912#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
913pub struct YamlOverrideTemplateEntry {
914    pub id: String,
915    pub name: String,
916    pub description: String,
917    /// Account type name from the IDL (overrides collection-level account_type)
918    #[serde(default)]
919    pub idl_account_name: Option<String>,
920    /// Properties with full metadata
921    #[serde(default)]
922    pub properties: Vec<YamlProperty>,
923    pub address: YamlAccountAddress,
924    /// Optional context/instructions specifically for LLMs using this template
925    #[serde(default)]
926    pub llm_context: Option<String>,
927}
928
929impl YamlOverrideTemplateCollection {
930    /// Convert collection to runtime OverrideTemplates with loaded IDL
931    pub fn to_override_templates(self, idl: Idl) -> Vec<OverrideTemplate> {
932        // Convert constants once for sharing
933        let constants: HashMap<String, ConstantDefinition> = self
934            .constants
935            .into_iter()
936            .map(|(k, v)| (k, v.into()))
937            .collect();
938
939        let default_account_type = self.account_type.clone().unwrap_or_default();
940
941        self.templates
942            .into_iter()
943            .map(|entry| OverrideTemplate {
944                id: entry.id,
945                name: entry.name,
946                description: entry.description,
947                protocol: self.protocol.clone(),
948                idl: idl.clone(),
949                address: entry.address.into(),
950                account_type: entry
951                    .idl_account_name
952                    .unwrap_or_else(|| default_account_type.clone()),
953                properties: entry.properties.into_iter().map(Into::into).collect(),
954                constants: constants.clone(),
955                tags: self.tags.clone(),
956                llm_context: entry.llm_context,
957            })
958            .collect()
959    }
960}
961
962/// YAML representation of an override template with embedded IDL
963/// Used for RPC methods where file access is not available
964#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
965pub struct YamlOverrideTemplate {
966    pub id: String,
967    pub name: String,
968    pub description: String,
969    pub protocol: String,
970    pub version: String,
971    pub account_type: String,
972    pub idl: Idl,
973    pub address: YamlAccountAddress,
974    #[serde(default)]
975    pub properties: Vec<YamlProperty>,
976    #[serde(default)]
977    pub constants: HashMap<String, YamlConstantDefinition>,
978    #[serde(default)]
979    pub tags: Vec<String>,
980    /// Optional context/instructions specifically for LLMs using this template
981    #[serde(default)]
982    pub llm_context: Option<String>,
983}
984
985impl YamlOverrideTemplate {
986    /// Convert to runtime OverrideTemplate
987    pub fn to_override_template(self) -> OverrideTemplate {
988        OverrideTemplate {
989            id: self.id,
990            name: self.name,
991            description: self.description,
992            protocol: self.protocol,
993            idl: self.idl,
994            address: self.address.into(),
995            account_type: self.account_type,
996            properties: self.properties.into_iter().map(Into::into).collect(),
997            constants: self
998                .constants
999                .into_iter()
1000                .map(|(k, v)| (k, v.into()))
1001                .collect(),
1002            tags: self.tags,
1003            llm_context: self.llm_context,
1004        }
1005    }
1006}
1007
1008/// YAML representation of account address
1009#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1010#[serde(tag = "type", rename_all = "lowercase")]
1011pub enum YamlAccountAddress {
1012    Pubkey {
1013        #[serde(default)]
1014        value: Option<String>,
1015    },
1016    Pda {
1017        program_id: String,
1018        seeds: Vec<YamlPdaSeed>,
1019    },
1020}
1021
1022impl From<YamlAccountAddress> for AccountAddress {
1023    fn from(yaml: YamlAccountAddress) -> Self {
1024        match yaml {
1025            YamlAccountAddress::Pubkey { value } => {
1026                AccountAddress::Pubkey(value.unwrap_or_default())
1027            }
1028            YamlAccountAddress::Pda { program_id, seeds } => AccountAddress::Pda {
1029                program_id,
1030                seeds: seeds.into_iter().map(|s| s.into()).collect(),
1031            },
1032        }
1033    }
1034}
1035
1036/// YAML representation of PDA seeds
1037#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1038#[serde(tag = "type", rename_all = "snake_case")]
1039pub enum YamlPdaSeed {
1040    String {
1041        value: String,
1042    },
1043    Bytes {
1044        value: Vec<u8>,
1045    },
1046    Pubkey {
1047        value: String,
1048    },
1049    PropertyRef {
1050        value: String,
1051    },
1052    /// A u16 value converted to big-endian bytes
1053    U16Be {
1054        value: u16,
1055    },
1056    /// Reference to a property that should be converted to u16 big-endian bytes
1057    U16BeRef {
1058        value: String,
1059    },
1060    /// A u16 value converted to little-endian bytes (useful for Pyth shard IDs)
1061    U16Le {
1062        value: u16,
1063    },
1064    /// Reference to a property that's a 32-byte hex string (e.g., Pyth feed ID)
1065    Bytes32Ref {
1066        value: String,
1067    },
1068    /// A nested PDA derivation - derives a PDA from inner seeds and uses it as the seed
1069    DerivedPda {
1070        program_id: String,
1071        seeds: Vec<YamlPdaSeed>,
1072    },
1073}
1074
1075impl From<YamlPdaSeed> for PdaSeed {
1076    fn from(yaml: YamlPdaSeed) -> Self {
1077        match yaml {
1078            YamlPdaSeed::String { value } => PdaSeed::String(value),
1079            YamlPdaSeed::Bytes { value } => PdaSeed::Bytes(value),
1080            YamlPdaSeed::Pubkey { value } => PdaSeed::Pubkey(value),
1081            YamlPdaSeed::PropertyRef { value } => PdaSeed::PropertyRef(value),
1082            YamlPdaSeed::U16Be { value } => PdaSeed::U16Be(value),
1083            YamlPdaSeed::U16BeRef { value } => PdaSeed::U16BeRef(value),
1084            YamlPdaSeed::U16Le { value } => PdaSeed::U16Le(value),
1085            YamlPdaSeed::Bytes32Ref { value } => PdaSeed::Bytes32Ref(value),
1086            YamlPdaSeed::DerivedPda { program_id, seeds } => PdaSeed::DerivedPda {
1087                program_id,
1088                seeds: seeds.into_iter().map(|s| s.into()).collect(),
1089            },
1090        }
1091    }
1092}
1093
1094#[cfg(test)]
1095mod tests {
1096    use std::collections::HashMap;
1097
1098    use serde_json::json;
1099
1100    use super::PdaSeed;
1101
1102    #[test]
1103    fn u16_be_ref_rejects_out_of_range_values() {
1104        let seed = PdaSeed::U16BeRef("index".to_string());
1105        let values = HashMap::from([("index".to_string(), json!(70_000))]);
1106
1107        assert_eq!(seed.to_bytes(Some(&values)), None);
1108    }
1109
1110    #[test]
1111    fn u16_be_ref_encodes_in_range_values() {
1112        let seed = PdaSeed::U16BeRef("index".to_string());
1113        let values = HashMap::from([("index".to_string(), json!(513))]);
1114
1115        assert_eq!(seed.to_bytes(Some(&values)), Some(vec![2, 1]));
1116    }
1117}