From f13d10c643a0c6493ff48812a109da495d90ea77 Mon Sep 17 00:00:00 2001 From: MicaiahReid Date: Thu, 13 Nov 2025 10:24:22 -0500 Subject: [PATCH 1/5] feat: add surfnet_registerScenario RPC method for scenario registration with account overrides --- crates/types/src/rpc_endpoints.json | 42 ++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/crates/types/src/rpc_endpoints.json b/crates/types/src/rpc_endpoints.json index 88ace52f..17a33893 100644 --- a/crates/types/src/rpc_endpoints.json +++ b/crates/types/src/rpc_endpoints.json @@ -763,6 +763,46 @@ } ], "returns": "A `RpcResponse<()>` indicating whether the write was successful." + }, + { + "method": "surfnet_registerScenario", + "description": "A cheat code to register a scenario with account overrides.", + "params": [ + { + "name": "scenario", + "type": "object", + "description": "The scenario object containing account overrides.", + "schema": { + "id": "String (A unique identifier for the scenario)", + "name": "String (A Human-readable name for the scenario)", + "description": "String (A description of the scenario)", + "overrides": { + "type": "array", + "description": "Array of OverrideInstance objects, each containing account overrides.", + "items": { + "type": "object", + "description": "An OverrideInstance object.", + "schema": { + "id": "String (A unique identifier for this override instance)", + "template_id": "String (The identifier of the template this instance is based on)", + "values": "HashMap (A hash map of field paths to override values. The field paths are a flat map to the account key being updated, from the IDL, using dot-notation to provide the path to the key. For example, to override a field 'amount' in an account 'myAccount', the path would be 'myAccount.amount'.)", + "scenario_relative_slot": "u64 (Relative slot when this override should be applied (relative to scenario registration slot))", + "label": "Option (An optional label for this override instance)", + "enabled": "bool (Indicates whether this override instance is enabled)", + "fetch_before_use": "bool (Indicates whether to fetch the latest on-chain account data before applying overrides)", + "account": "String (The account this override instance is applied to)" + } + } + } + } + }, + { + "name": "slot", + "type": "Option", + "description": "The base slot from which relative slot offsets are calculated. If omitted, uses the current slot." + } + ], + "returns": "A `RpcResponse<()>` indicating whether the write was successful." } ] }, @@ -1444,4 +1484,4 @@ ] } ] -} \ No newline at end of file +} From 21c47786c8203009b7c424463971bf4635557297 Mon Sep 17 00:00:00 2001 From: MicaiahReid Date: Thu, 13 Nov 2025 14:50:56 -0500 Subject: [PATCH 2/5] feat: implement register_scenario RPC method for scenario creation with account overrides --- crates/mcp/Cargo.toml | 1 + crates/mcp/src/surfpool/mod.rs | 271 +++++++++++++++++++++++++++++++-- 2 files changed, 260 insertions(+), 12 deletions(-) diff --git a/crates/mcp/Cargo.toml b/crates/mcp/Cargo.toml index cedba610..4aa08b49 100644 --- a/crates/mcp/Cargo.toml +++ b/crates/mcp/Cargo.toml @@ -20,6 +20,7 @@ reqwest = { workspace = true } rmcp = { workspace = true, features = ["transport-io", "transport-sse-server", "server", "schemars"] } serde = { workspace = true } serde_json = { workspace = true } +serde_yaml = "0.9" solana-keypair = { workspace = true } solana-pubkey = { workspace = true } solana-signer = { workspace = true } diff --git a/crates/mcp/src/surfpool/mod.rs b/crates/mcp/src/surfpool/mod.rs index 64adb059..ba64307e 100644 --- a/crates/mcp/src/surfpool/mod.rs +++ b/crates/mcp/src/surfpool/mod.rs @@ -13,15 +13,19 @@ use rmcp::{ }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Value, json}; use set_token_account::{SeededAccount, SetAccountSuccess, SetTokenAccountsResponse}; use start_surfnet::StartSurfnetResponse; +use surfpool_types::YamlOverrideTemplateCollection; use crate::helpers::find_next_available_surfnet_port; mod set_token_account; mod start_surfnet; +pub const PYTH_V2_OVERRIDES_CONTENT: &str = + include_str!("../../../core/src/scenarios/protocols/pyth/v2/overrides.yaml"); + #[derive(Debug, Clone)] pub struct Surfpool { pub surfnets: Arc>>, @@ -175,6 +179,21 @@ impl SurfnetRpcCallResponse { } } +#[derive(Serialize, Deserialize, JsonSchema)] +pub struct RegisterScenarioResponse { + pub error: Option, +} +impl RegisterScenarioResponse { + pub fn success() -> Self { + Self { error: None } + } + pub fn error(message: String) -> Self { + Self { + error: Some(message), + } + } +} + #[tool(tool_box)] impl Surfpool { /// Returns a command to start a new local Solana network (surfnet). @@ -334,9 +353,15 @@ impl Surfpool { /// Calls any RPC method on a running surfnet instance. /// This generic method allows calling any of the available surfnet cheatcode RPC methods. /// The LLM will interpret user requests and determine which method to call with appropriate parameters. - #[tool( - description = "Calls any RPC method on a running surfnet instance. This is a generic method that can invoke any surfnet RPC method. The LLM should interpret user requests and determine the appropriate method and parameters to call. To retrieve the list of RPC endpoints available check the resource str:///rpc_endpoints" - )] + #[tool(description = r#" + Calls any RPC method on a running surfnet instance. + This is a generic method that can invoke any surfnet RPC method. + The LLM should interpret user requests and determine the appropriate method and parameters to call. To retrieve the list of RPC endpoints available check the resource str:///rpc_endpoints + + IMPORTANT: + - If a user asks to create a scenario, DO NOT USE this method. Instead, use the dedicated `register_scenario` tool. + - There is NO RPC method called `surfnet_getOverrideTemplates`. Override templates are ONLY available via the MCP resource str:///override_templates (use read_resource, not this RPC tool). + "#)] pub fn call_surfnet_rpc( &self, #[tool(param)] @@ -417,6 +442,203 @@ impl Surfpool { Json(SurfnetRpcCallResponse::success(method, result)) } + + #[tool(description = r#" + This tool creates a scenario that overrides account data on a running surfnet instance. + It will pass the provided parameters to another service that handles the `surfnet_registerScenario` RPC method. + + CRITICAL PREREQUISITE: You MUST use read_resource to fetch str:///override_templates BEFORE calling this tool. + If read_resource is not available, use the tool `get_override_templates` to get the override templates. + This resource contains all available override templates, their field paths (for the `values` map), and account addresses. + DO NOT attempt to call any RPC method to get templates - they are ONLY available via the MCP resource system. + + The first argument should be the port of the running surfnet instance (e.g., 8899, 18899, 28899, etc.). + The second argument should be the parameters to pass to the `surfnet_registerScenario` RPC method. This should be a JSON object containing: + - `scenario`: The Scenario object containing: + - `id`: Unique identifier for the scenario + - `name`: Human-readable name + - `description`: Description of the scenario + - `overrides`: Array of OverrideInstance objects, each containing: + - `id`: Unique identifier for this override instance + - `templateId`: Reference to the override template + - `values`: HashMap of field paths to override values (flat key-value map with dot notation, e.g., "price_message.price_value") + - `scenarioRelativeSlot`: The relative slot offset (from base slot) when this override should be applied + - `label`: Optional label for this override + - `enabled`: Whether this override is active + - `fetchBeforeUse`: If true, fetch fresh account data just before transaction execution (useful for price feeds, oracle updates, and dynamic balances) + - `account`: Account address (either `{ "pubkey": "..." }` or `{ "pda": { "programId": "...", "seeds": [...] } }`) + - `tags`: Array of tags for categorization + - `slot` (optional): The base slot from which relative slot offsets are calculated. If omitted, uses the current slot. + + REMINDER: Valid keys for the scenario's override's values field and the account addresses MUST come from str:///override_templates. + You MUST fetch this resource using read_resource before constructing the scenario parameters. + + For example, if the str:///override_templates resource returns: + ```json + { + "pyth_v2": { + "protocol": "Pyth", + "version": "v2", + "idl_file_path": "pyth_price_store.json", + "tags": [ + "oracle", + "price-feed", + "defi" + ], + "templates": [ + { + "id": "pyth-sol-usd-v2", + "name": "Override SOL/USD Price Feed", + "description": "Override Pyth SOL/USD price feed with custom price data", + "idl_account_name": "PriceUpdateV2", + "properties": [ + "price_message.price", + "price_message.publish_time" + ], + "address": { + "type": "pubkey", + "value": "7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE" + } + } + ] + } + } + ``` + + And the user asks to create a scenario that sets the SOL/USD price to $150.00, the `params` to pass to this method would be: + ```json + { + "id": "scenario-001", + "name": "Set SOL/USD Price to $150", + "description": "A scenario that sets the SOL/USD price feed to $150.00 at slot 5000", + "overrides": [ + { + "id": "override-001", + "templateId": "pyth-sol-usd-v2", + "values": { + "price_message.price": 150000000, // Pyth price is in 10^-6 format + "price_message.publish_time": 1625247600 // Example publish time + }, + "scenarioRelativeSlot": 1, + "label": "Set SOL/USD to $150", + "enabled": true, + "fetchBeforeUse": false, + "account": { + "pubkey": "7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE" + } + } + ], + "tags": ["test", "price-feed"] + } + ``` + + "#)] + pub fn register_scenario( + &self, + #[tool(param)] + #[schemars( + description = r#"The parameters to pass to the `surfnet_registerScenario` RPC method. This should contain a `Scenario` object containing a list of `OverrideInstance`s. + Example: + ```json + [ + { + "id": "scenario-001", + "name": "Set SOL/USD Price to $150", + "description": "A scenario that sets the SOL/USD price feed to $150.00 at slot 5000", + "overrides": [ + { + "id": "override-001", + "templateId": "pyth-sol-usd-v2", + "values": { + "price_message.price": 150000000, // Pyth price is in 10^-6 format + "price_message.publish_time": 1625247600 // Example publish time + }, + "scenarioRelativeSlot": 1, + "label": "Set SOL/USD to $150", + "enabled": true, + "fetchBeforeUse": false, + "account": { + "pubkey": "7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE" + } + } + ], + "tags": ["test", "price-feed"] + } + ] + ``` + "# + )] + params: Vec, + ) -> Json { + let load_scenarios_endpoint = format!("http://127.0.0.1:{}/v1/scenarios", 18488); + // let params: Vec = params.into_iter().map(Into::into).collect(); + let payload = serde_json::json!(params); + + // Make the RPC request to the surfnet RPC endpoint + let client = reqwest::blocking::Client::new(); + let response = match client + .post(&load_scenarios_endpoint) + .header("Content-Type", "application/json") + .json(&payload) + .send() + { + Ok(resp) => resp, + Err(e) => { + return Json(RegisterScenarioResponse::error(format!( + "Failed to load scenarios at {}: {}", + load_scenarios_endpoint, e + ))); + } + }; + + let response_text = match response.text() { + Ok(text) => text, + Err(e) => { + return Json(RegisterScenarioResponse::error(format!( + "Failed to read response text: {}", + e + ))); + } + }; + + let rpc_response: serde_json::Value = match serde_json::from_str(&response_text) { + Ok(json) => json, + Err(e) => { + return Json(RegisterScenarioResponse::error(format!( + "Failed to parse JSON response: {}. Response: {}", + e, response_text + ))); + } + }; + + if let Some(error) = rpc_response.get("error") { + return Json(RegisterScenarioResponse::error(format!( + "RPC error: {}", + error + ))); + } + + // Extract the result + let result = rpc_response + .get("result") + .unwrap_or(&serde_json::Value::Null) + .clone(); + + Json(RegisterScenarioResponse::success()) + } + + #[tool( + description = "Fetches the override templates resource (equivalent to str:///override_templates)." + )] + pub fn get_override_templates(&self) -> Json { + let pyth_v2_json_value = + serde_yaml::from_str::(PYTH_V2_OVERRIDES_CONTENT) + .expect("Expected Pyth overrides file to be deserializable"); + + Json(json!({ + "pyth_v2": pyth_v2_json_value, + })) + } } #[tool(tool_box)] @@ -441,14 +663,22 @@ impl ServerHandler for Surfpool { _: RequestContext, ) -> Result { Ok(ListResourcesResult { - resources: vec![RawResource { - uri: "str:///rpc_endpoints".to_string(), - name: "List of available RPC endpoints".to_string(), - description: Some("A json file containing all the RPC methods and the parameters available for being able to handle any RPC call with the tool call_surfnet_rpc".to_string()), - mime_type: Some("application/json".to_string()), - size: None, - } - .no_annotation()], + resources: vec![ + RawResource { + uri: "str:///rpc_endpoints".to_string(), + name: "List of available RPC endpoints".to_string(), + description: Some("A json file containing all the RPC methods and the parameters available for being able to handle any RPC call with the tool call_surfnet_rpc".to_string()), + mime_type: Some("application/json".to_string()), + size: None, + }.no_annotation(), + RawResource { + uri: "str:///override_templates".to_string(), + name: "List of override templates".to_string(), + description: Some("A json file containing all the override templates available for scenario registration with account overrides".to_string()), + mime_type: Some("application/json".to_string()), + size: None, + }.no_annotation() + ], next_cursor: None, }) } @@ -468,6 +698,23 @@ impl ServerHandler for Surfpool { }], }) } + "str:///override_templates" => { + let pyth_v2_json_value = serde_yaml::from_str::( + PYTH_V2_OVERRIDES_CONTENT, + ) + .expect("Expected Pyth overrides file to be deserializable"); + let joined_value = json!({ + "pyth_v2": pyth_v2_json_value, + }); + + Ok(ReadResourceResult { + contents: vec![ResourceContents::TextResourceContents { + uri, + mime_type: Some("application/json".to_string()), + text: joined_value.to_string(), + }], + }) + } _ => Err(McpError::resource_not_found( "resource_not_found", Some(serde_json::json!({ From 8c919e79dad44966e4a463d6c771994c2387f963 Mon Sep 17 00:00:00 2001 From: MicaiahReid Date: Thu, 13 Nov 2025 14:51:07 -0500 Subject: [PATCH 3/5] feat: update surfpool packages to version 0.12.0 and add scenario handling endpoints --- Cargo.lock | 24 ++++++++------- Cargo.toml | 6 ++-- crates/cli/src/http/mod.rs | 60 ++++++++++++++++++++++++++++++++++++-- 3 files changed, 74 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bcd9b68f..8fbf3367 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12013,7 +12013,7 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "surfpool-cli" -version = "0.11.2" +version = "0.12.0" dependencies = [ "actix-cors", "actix-web", @@ -12070,7 +12070,7 @@ dependencies = [ [[package]] name = "surfpool-core" -version = "0.11.2" +version = "0.12.0" dependencies = [ "agave-feature-set", "agave-geyser-plugin-interface", @@ -12155,7 +12155,7 @@ dependencies = [ [[package]] name = "surfpool-db" -version = "0.11.2" +version = "0.12.0" dependencies = [ "diesel", "diesel-dynamic-schema", @@ -12168,7 +12168,7 @@ dependencies = [ [[package]] name = "surfpool-gql" -version = "0.11.2" +version = "0.12.0" dependencies = [ "base64 0.22.1", "blake3", @@ -12189,7 +12189,7 @@ dependencies = [ [[package]] name = "surfpool-mcp" -version = "0.11.2" +version = "0.12.0" dependencies = [ "bs58", "crossbeam-channel", @@ -12198,6 +12198,7 @@ dependencies = [ "rmcp", "serde", "serde_json", + "serde_yaml", "solana-keypair", "solana-pubkey 3.0.0", "solana-signer", @@ -12233,7 +12234,7 @@ dependencies = [ [[package]] name = "surfpool-subgraph" -version = "0.11.2" +version = "0.12.0" dependencies = [ "agave-geyser-plugin-interface", "ipc-channel", @@ -12248,7 +12249,7 @@ dependencies = [ [[package]] name = "surfpool-types" -version = "0.11.2" +version = "0.12.0" dependencies = [ "anchor-lang-idl", "blake3", @@ -13054,9 +13055,9 @@ dependencies = [ [[package]] name = "txtx-addon-network-svm" -version = "0.3.16" +version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e956b4c0ed49913bb4a08b3e6062a68f9aa71dab66c06e6f3abe88c9878fcb9d" +checksum = "23a089a07536356a6ae82e87c0269497008043950939ea2743b9d97b0d3615e4" dependencies = [ "async-recursion", "bincode", @@ -13065,6 +13066,7 @@ dependencies = [ "convert_case 0.6.0", "kaigan", "lazy_static", + "log 0.4.28", "serde", "serde_derive", "serde_json", @@ -13098,9 +13100,9 @@ dependencies = [ [[package]] name = "txtx-addon-network-svm-types" -version = "0.3.15" +version = "0.3.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30eaa1527896106eea3fce87bf1aac932c31c45ff94b544c0eabbec2f63ebd01" +checksum = "0df4bf99587f0513d9af93899afcbd99a43e7826dc9d568a6f2646fca7493575" dependencies = [ "anchor-lang-idl", "borsh 1.5.7", diff --git a/Cargo.toml b/Cargo.toml index 9af6302a..e28f0e69 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace.package] -version = "0.11.2" +version = "0.12.0" edition = "2024" description = "Surfpool is where developers start their Solana journey." license = "Apache-2.0" @@ -160,8 +160,8 @@ surfpool-subgraph = { path = "crates/subgraph", default-features = false } surfpool-types = { path = "crates/types", default-features = false } txtx-addon-kit = "0.4.10" -txtx-addon-network-svm = { version = "0.3.16" } -txtx-addon-network-svm-types = { version = "0.3.15" } +txtx-addon-network-svm = { version = "0.3.17" } +txtx-addon-network-svm-types = { version = "0.3.16" } txtx-cloud = { version = "0.1.13", features = [ "clap", "toml", diff --git a/crates/cli/src/http/mod.rs b/crates/cli/src/http/mod.rs index a2c48ca2..d70a129d 100644 --- a/crates/cli/src/http/mod.rs +++ b/crates/cli/src/http/mod.rs @@ -1,6 +1,9 @@ #![allow(unused_imports, unused_variables)] use std::{ - collections::HashMap, error::Error as StdError, sync::RwLock, thread::JoinHandle, + collections::HashMap, + error::Error as StdError, + sync::{Arc, RwLock}, + thread::JoinHandle, time::Duration, }; @@ -9,7 +12,7 @@ use actix_web::{ App, Error, HttpRequest, HttpResponse, HttpServer, Responder, dev::ServerHandle, http::header::{self}, - middleware, + middleware, post, web::{self, Data, route}, }; use convert_case::{Case, Casing}; @@ -19,6 +22,7 @@ use juniper_graphql_ws::ConnectionConfig; use log::{debug, error, info, trace, warn}; #[cfg(feature = "explorer")] use rust_embed::RustEmbed; +use serde::{Deserialize, Serialize}; use surfpool_core::scenarios::TemplateRegistry; use surfpool_gql::{ DynamicSchema, @@ -68,6 +72,7 @@ pub async fn start_subgraph_and_explorer_server( // Initialize template registry and load templates let template_registry_wrapped = Data::new(RwLock::new(TemplateRegistry::new())); + let loaded_scenarios = Data::new(RwLock::new(LoadedScenarios::new())); let subgraph_handle = start_subgraph_runloop( subgraph_events_tx, @@ -86,6 +91,7 @@ pub async fn start_subgraph_and_explorer_server( .app_data(config_wrapped.clone()) .app_data(collections_metadata_lookup_wrapped.clone()) .app_data(template_registry_wrapped.clone()) + .app_data(loaded_scenarios.clone()) .wrap( Cors::default() .allow_any_origin() @@ -100,6 +106,8 @@ pub async fn start_subgraph_and_explorer_server( .service(get_config) .service(get_indexers) .service(get_scenario_templates) + .service(post_scenarios) + .service(get_scenarios) .service( web::scope("/workspace") .route("/v1/indexers", web::post().to(post_graphql)) @@ -109,6 +117,7 @@ pub async fn start_subgraph_and_explorer_server( ); if enable_studio { + app = app.app_data(Arc::new(RwLock::new(LoadedScenarios::new()))); app = app.service(serve_studio_static_files); } @@ -191,6 +200,53 @@ async fn get_scenario_templates( .body(response)) } +#[derive(Debug, Serialize, Deserialize)] +pub struct LoadedScenarios { + pub scenarios: Vec, +} +impl LoadedScenarios { + pub fn new() -> Self { + Self { + scenarios: Vec::new(), + } + } +} + +#[post("/v1/scenarios")] +async fn post_scenarios( + req: HttpRequest, + payload: web::Payload, + data: Data>, +) -> Result { + let scenario = serde_json::from_slice::( + &payload + .to_bytes() + .await + .map_err(|_| actix_web::error::ErrorBadRequest("Failed to read request payload"))?, + ) + .map_err(|_| actix_web::error::ErrorBadRequest("Failed to parse JSON"))?; + + let mut loaded_scenarios = data + .write() + .map_err(|_| actix_web::error::ErrorInternalServerError("Failed to acquire write lock"))?; + loaded_scenarios.scenarios.push(scenario); + Ok(HttpResponse::Ok().body("Scenario loaded")) +} + +#[actix_web::get("/v1/scenarios")] +async fn get_scenarios(data: Data>) -> Result { + let loaded_scenarios = data + .read() + .map_err(|_| actix_web::error::ErrorInternalServerError("Failed to acquire read lock"))?; + let response = serde_json::to_string(&loaded_scenarios.scenarios).map_err(|_| { + actix_web::error::ErrorInternalServerError("Failed to serialize loaded scenarios") + })?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(response)) +} + #[allow(dead_code)] #[cfg(not(feature = "explorer"))] fn handle_embedded_file(_path: &str) -> HttpResponse { From ec9652cce4dc73f618de7084d402853dd09ef1d3 Mon Sep 17 00:00:00 2001 From: Ludo Galabru Date: Thu, 13 Nov 2025 22:02:13 -0500 Subject: [PATCH 4/5] feat: mcp improvements --- Cargo.lock | 2 + Cargo.toml | 2 + crates/cli/src/http/mod.rs | 88 ++++++++++++--- crates/mcp/src/surfpool/mod.rs | 201 +++++++++++---------------------- crates/types/Cargo.toml | 2 + crates/types/src/scenarios.rs | 28 ++++- 6 files changed, 162 insertions(+), 161 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8fbf3367..70579e4a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12256,6 +12256,8 @@ dependencies = [ "chrono", "crossbeam-channel", "once_cell", + "schemars 0.8.22", + "schemars_derive", "serde", "serde_json", "serde_with", diff --git a/Cargo.toml b/Cargo.toml index e28f0e69..85a236fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -96,6 +96,8 @@ ratatui = { version = "0.29.0", features = [ reqwest = { version = "0.12.23", default-features = false } rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", rev = "ff71a526156e6c9409c450f71eccd6aced9bc339", package = "rmcp" } rust-embed = "8.2.0" +schemars = { version = "0.8.22" } +schemars_derive = { version = "0.8.22" } serde = { version = "1.0.226", default-features = false } serde_derive = { version = "1.0.226", default-features = false } # must match the serde version, see https://github.com/serde-rs/serde/issues/2584#issuecomment-1685252251 serde_json = { version = "1.0.135", default-features = false } diff --git a/crates/cli/src/http/mod.rs b/crates/cli/src/http/mod.rs index d70a129d..994086db 100644 --- a/crates/cli/src/http/mod.rs +++ b/crates/cli/src/http/mod.rs @@ -33,8 +33,8 @@ use surfpool_gql::{ }; use surfpool_studio_ui::serve_studio_static_files; use surfpool_types::{ - DataIndexingCommand, OverrideTemplate, SanitizedConfig, SubgraphCommand, SubgraphEvent, - SurfpoolConfig, + DataIndexingCommand, OverrideTemplate, SanitizedConfig, Scenario, SubgraphCommand, + SubgraphEvent, SurfpoolConfig, }; use txtx_core::kit::types::types::Value; use txtx_gql::kit::uuid::Uuid; @@ -95,9 +95,8 @@ pub async fn start_subgraph_and_explorer_server( .wrap( Cors::default() .allow_any_origin() - .allowed_methods(vec!["POST", "GET", "OPTIONS", "DELETE"]) - .allowed_headers(vec![header::AUTHORIZATION, header::ACCEPT]) - .allowed_header(header::CONTENT_TYPE) + .allow_any_method() + .allow_any_header() .supports_credentials() .max_age(3600), ) @@ -108,6 +107,8 @@ pub async fn start_subgraph_and_explorer_server( .service(get_scenario_templates) .service(post_scenarios) .service(get_scenarios) + .service(delete_scenario) + .service(patch_scenario) .service( web::scope("/workspace") .route("/v1/indexers", web::post().to(post_graphql)) @@ -202,7 +203,7 @@ async fn get_scenario_templates( #[derive(Debug, Serialize, Deserialize)] pub struct LoadedScenarios { - pub scenarios: Vec, + pub scenarios: Vec, } impl LoadedScenarios { pub fn new() -> Self { @@ -215,22 +216,19 @@ impl LoadedScenarios { #[post("/v1/scenarios")] async fn post_scenarios( req: HttpRequest, - payload: web::Payload, + scenario: web::Json, data: Data>, ) -> Result { - let scenario = serde_json::from_slice::( - &payload - .to_bytes() - .await - .map_err(|_| actix_web::error::ErrorBadRequest("Failed to read request payload"))?, - ) - .map_err(|_| actix_web::error::ErrorBadRequest("Failed to parse JSON"))?; - let mut loaded_scenarios = data .write() .map_err(|_| actix_web::error::ErrorInternalServerError("Failed to acquire write lock"))?; - loaded_scenarios.scenarios.push(scenario); - Ok(HttpResponse::Ok().body("Scenario loaded")) + let scenario_data = scenario.into_inner(); + let scenario_id = scenario_data.id.clone(); + loaded_scenarios.scenarios.push(scenario_data); + let response = serde_json::json!({"id": scenario_id}); + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(response.to_string())) } #[actix_web::get("/v1/scenarios")] @@ -247,6 +245,62 @@ async fn get_scenarios(data: Data>) -> Result, + data: Data>, +) -> Result { + let scenario_id = path.into_inner(); + let mut loaded_scenarios = data + .write() + .map_err(|_| actix_web::error::ErrorInternalServerError("Failed to acquire write lock"))?; + + let initial_len = loaded_scenarios.scenarios.len(); + loaded_scenarios.scenarios.retain(|s| s.id != scenario_id); + + if loaded_scenarios.scenarios.len() == initial_len { + return Ok( + HttpResponse::NotFound().body(format!("Scenario with id '{}' not found", scenario_id)) + ); + } + + Ok(HttpResponse::Ok().body(format!("Scenario '{}' deleted", scenario_id))) +} + +#[actix_web::patch("/v1/scenarios/{id}")] +async fn patch_scenario( + path: web::Path, + scenario: web::Json, + data: Data>, +) -> Result { + let scenario_id = path.into_inner(); + let mut loaded_scenarios = data + .write() + .map_err(|_| actix_web::error::ErrorInternalServerError("Failed to acquire write lock"))?; + + let scenario_index = loaded_scenarios + .scenarios + .iter() + .position(|s| s.id == scenario_id); + + match scenario_index { + Some(index) => { + loaded_scenarios.scenarios[index] = scenario.into_inner(); + let response = serde_json::json!({"id": scenario_id}); + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(response.to_string())) + } + None => { + loaded_scenarios.scenarios.push(scenario.into_inner()); + let response = serde_json::json!({"id": scenario_id}); + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(response.to_string())) + } + } +} + #[allow(dead_code)] #[cfg(not(feature = "explorer"))] fn handle_embedded_file(_path: &str) -> HttpResponse { diff --git a/crates/mcp/src/surfpool/mod.rs b/crates/mcp/src/surfpool/mod.rs index ba64307e..dc018c39 100644 --- a/crates/mcp/src/surfpool/mod.rs +++ b/crates/mcp/src/surfpool/mod.rs @@ -16,7 +16,10 @@ use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use set_token_account::{SeededAccount, SetAccountSuccess, SetTokenAccountsResponse}; use start_surfnet::StartSurfnetResponse; -use surfpool_types::YamlOverrideTemplateCollection; +use surfpool_core::scenarios::TemplateRegistry; +use surfpool_types::{ + CHANGE_TO_DEFAULT_STUDIO_PORT_ONCE_SUPERVISOR_MERGED, Scenario, YamlOverrideTemplateCollection, +}; use crate::helpers::find_next_available_surfnet_port; @@ -29,12 +32,14 @@ pub const PYTH_V2_OVERRIDES_CONTENT: &str = #[derive(Debug, Clone)] pub struct Surfpool { pub surfnets: Arc>>, + pub template_registry: Arc>, } impl Surfpool { pub fn new() -> Self { Self { surfnets: Arc::new(RwLock::new(HashMap::new())), + template_registry: Arc::new(RwLock::new(TemplateRegistry::new())), } } } @@ -182,14 +187,19 @@ impl SurfnetRpcCallResponse { #[derive(Serialize, Deserialize, JsonSchema)] pub struct RegisterScenarioResponse { pub error: Option, + pub url: Option, } impl RegisterScenarioResponse { - pub fn success() -> Self { - Self { error: None } + pub fn success(url: String) -> Self { + Self { + error: None, + url: Some(url), + } } pub fn error(message: String) -> Self { Self { error: Some(message), + url: None, } } } @@ -201,7 +211,7 @@ impl Surfpool { #[tool(description = "Starts a new local Solana network (surfnet). This tool returns a shell command that the AI must execute in a new terminal. This is the default behavior (`run_as_subprocess: false`). - If the user explicitly asks to run surfnet as a background process, + If the user explicitly asks to run surfnet as a background process, set `run_as_subprocess: true` to start it directly. When in doubt, return the command for execution.")] pub fn start_surfnet( &self, @@ -354,12 +364,12 @@ impl Surfpool { /// This generic method allows calling any of the available surfnet cheatcode RPC methods. /// The LLM will interpret user requests and determine which method to call with appropriate parameters. #[tool(description = r#" - Calls any RPC method on a running surfnet instance. - This is a generic method that can invoke any surfnet RPC method. + Calls any RPC method on a running surfnet instance. + This is a generic method that can invoke any surfnet RPC method. The LLM should interpret user requests and determine the appropriate method and parameters to call. To retrieve the list of RPC endpoints available check the resource str:///rpc_endpoints - - IMPORTANT: - - If a user asks to create a scenario, DO NOT USE this method. Instead, use the dedicated `register_scenario` tool. + + IMPORTANT: + - If a user asks to create a scenario, DO NOT USE this method. Instead, use the dedicated `create_scenario` tool. - There is NO RPC method called `surfnet_getOverrideTemplates`. Override templates are ONLY available via the MCP resource str:///override_templates (use read_resource, not this RPC tool). "#)] pub fn call_surfnet_rpc( @@ -372,8 +382,8 @@ impl Surfpool { #[tool(param)] #[schemars( description = "The RPC method name to call,for example: 'sendTransaction', 'simulateTransaction', 'getAccountInfo', 'getBalance', 'getTokenAccountBalance', 'getTokenSupply', 'getProgramAccounts', 'getTokenAccountsByOwner', 'getSlot', - 'getEpochInfo', 'requestAirdrop', 'surfnet_setAccount', 'getHealth', 'getTokenAccountsByOwner', 'getTokenAccountsByDelegate', - 'getTokenAccountsByDelegateAndMint', 'getTokenAccountsByDelegateAndMintAndOwner', 'getTokenAccountsByDelegateAndMintAndOwnerAndProgramId', 'getTokenAccountsByDelegateAndMintAndOwnerAndProgramIdAndOwner', surfnet_getProfileResults, etc. + 'getEpochInfo', 'requestAirdrop', 'surfnet_setAccount', 'getHealth', 'getTokenAccountsByOwner', 'getTokenAccountsByDelegate', + 'getTokenAccountsByDelegateAndMint', 'getTokenAccountsByDelegateAndMintAndOwner', 'getTokenAccountsByDelegateAndMintAndOwnerAndProgramId', 'getTokenAccountsByDelegateAndMintAndOwnerAndProgramIdAndOwner', surfnet_getProfileResults, etc. A list of all the RPC methods available can be found at str:///rpc_endpoints" )] method: String, @@ -444,135 +454,46 @@ impl Surfpool { } #[tool(description = r#" - This tool creates a scenario that overrides account data on a running surfnet instance. - It will pass the provided parameters to another service that handles the `surfnet_registerScenario` RPC method. + This tool creates a Scenario. A scenario is a list of state fragments that are used for overriding account data. + The tool returns the URL opening Surfpool Studio that the user can open for testing his scenario. - CRITICAL PREREQUISITE: You MUST use read_resource to fetch str:///override_templates BEFORE calling this tool. + CRITICAL PREREQUISITE: You MUST use read_resource to fetch the MCP resource "str:///override_templates" BEFORE calling this tool. If read_resource is not available, use the tool `get_override_templates` to get the override templates. This resource contains all available override templates, their field paths (for the `values` map), and account addresses. DO NOT attempt to call any RPC method to get templates - they are ONLY available via the MCP resource system. - - The first argument should be the port of the running surfnet instance (e.g., 8899, 18899, 28899, etc.). - The second argument should be the parameters to pass to the `surfnet_registerScenario` RPC method. This should be a JSON object containing: - - `scenario`: The Scenario object containing: - - `id`: Unique identifier for the scenario - - `name`: Human-readable name - - `description`: Description of the scenario - - `overrides`: Array of OverrideInstance objects, each containing: - - `id`: Unique identifier for this override instance - - `templateId`: Reference to the override template - - `values`: HashMap of field paths to override values (flat key-value map with dot notation, e.g., "price_message.price_value") - - `scenarioRelativeSlot`: The relative slot offset (from base slot) when this override should be applied - - `label`: Optional label for this override - - `enabled`: Whether this override is active - - `fetchBeforeUse`: If true, fetch fresh account data just before transaction execution (useful for price feeds, oracle updates, and dynamic balances) - - `account`: Account address (either `{ "pubkey": "..." }` or `{ "pda": { "programId": "...", "seeds": [...] } }`) - - `tags`: Array of tags for categorization - - `slot` (optional): The base slot from which relative slot offsets are calculated. If omitted, uses the current slot. + + You should take the templates and build the scenario structure out of it. The only data you can customize are: + - `id`: uuid v4 + - `name`: Human-readable name + - `description`: Description of the scenario + - `overrides[*].id`: uuid v4 + - `overrides[*].label`: Human-readable label for the override + - `overrides[*].scenarioRelativeSlot`: The relative slot offset (from base slot) when this override should be applied + - `overrides[*].values`: The values that should be overridden + All the other fields should be coming from the override templates. REMINDER: Valid keys for the scenario's override's values field and the account addresses MUST come from str:///override_templates. You MUST fetch this resource using read_resource before constructing the scenario parameters. For example, if the str:///override_templates resource returns: ```json - { - "pyth_v2": { - "protocol": "Pyth", - "version": "v2", - "idl_file_path": "pyth_price_store.json", - "tags": [ - "oracle", - "price-feed", - "defi" - ], - "templates": [ - { - "id": "pyth-sol-usd-v2", - "name": "Override SOL/USD Price Feed", - "description": "Override Pyth SOL/USD price feed with custom price data", - "idl_account_name": "PriceUpdateV2", - "properties": [ - "price_message.price", - "price_message.publish_time" - ], - "address": { - "type": "pubkey", - "value": "7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE" - } - } - ] - } - } + [{"id":"pyth-btc-usd-v2","name":"Override BTC/USD Price Feed","description":"Override Pyth BTC/USD price feed with custom price data","protocol":"Pyth","idl":{"address":"rec5EKMGg6MxZYaMdyBfgwp4d5rB9T1VQH5pJv5LtFJ","metadata":{"name":"price_feed","version":"0.1.0","spec":"0.1.0","description":"Created with Anchor"},"instructions":[],"accounts":[{"name":"PriceUpdateV2","discriminator":[34,241,35,99,157,126,244,205]}],"types":[{"name":"PriceFeedMessage","type":{"kind":"struct","fields":[{"name":"feed_id","docs":[" "],"type":{"array":["u8",32]}},{"name":"price","type":"i64"},{"name":"conf","type":"u64"},{"name":"exponent","type":"i32"},{"name":"publish_time","docs":["The timestamp of this price update in seconds"],"type":"i64"},{"name":"prev_publish_time","docs":[" ",""," "],"type":"i64"},{"name":"ema_price","type":"i64"},{"name":"ema_conf","type":"u64"}]}},{"name":"PriceUpdateV2","type":{"kind":"struct","fields":[{"name":"write_authority","type":"pubkey"},{"name":"verification_level","type":{"defined":{"name":"VerificationLevel"}}},{"name":"price_message","type":{"defined":{"name":"PriceFeedMessage"}}},{"name":"posted_slot","type":"u64"}]}},{"name":"VerificationLevel","type":{"kind":"enum","variants":[{"name":"Partial","fields":[{"name":"num_signatures","type":"u8"}]},{"name":"Full"}]}}]},"address":{"pubkey":"4cSM2e6rvbGQUFiJbqytoVMi5GgghSMr8LwVrT9VPSPo"},"accountType":"PriceUpdateV2","properties":["price_message.price","price_message.publish_time"],"tags":["oracle","price-feed","defi"]},{"id":"pyth-eth-btc-v2","name":"Override ETH/BTC Price Feed","description":"Override Pyth ETH/BTC price feed with custom price data","protocol":"Pyth","idl":{"address":"rec5EKMGg6MxZYaMdyBfgwp4d5rB9T1VQH5pJv5LtFJ","metadata":{"name":"price_feed","version":"0.1.0","spec":"0.1.0","description":"Created with Anchor"},"instructions":[],"accounts":[{"name":"PriceUpdateV2","discriminator":[34,241,35,99,157,126,244,205]}],"types":[{"name":"PriceFeedMessage","type":{"kind":"struct","fields":[{"name":"feed_id","docs":[" "],"type":{"array":["u8",32]}},{"name":"price","type":"i64"},{"name":"conf","type":"u64"},{"name":"exponent","type":"i32"},{"name":"publish_time","docs":["The timestamp of this price update in seconds"],"type":"i64"},{"name":"prev_publish_time","docs":[" ",""," "],"type":"i64"},{"name":"ema_price","type":"i64"},{"name":"ema_conf","type":"u64"}]}},{"name":"PriceUpdateV2","type":{"kind":"struct","fields":[{"name":"write_authority","type":"pubkey"},{"name":"verification_level","type":{"defined":{"name":"VerificationLevel"}}},{"name":"price_message","type":{"defined":{"name":"PriceFeedMessage"}}},{"name":"posted_slot","type":"u64"}]}},{"name":"VerificationLevel","type":{"kind":"enum","variants":[{"name":"Partial","fields":[{"name":"num_signatures","type":"u8"}]},{"name":"Full"}]}}]},"address":{"pubkey":"5JwbqPPMNpzE2jVAdobWo6m5gkhsDhRdGBo3FYbSfmaK"},"accountType":"PriceUpdateV2","properties":["price_message.price","price_message.publish_time"],"tags":["oracle","price-feed","defi"]},{"id":"pyth-eth-usd-v2","name":"Override ETH/USD Price Feed","description":"Override Pyth ETH/USD price feed with custom price data","protocol":"Pyth","idl":{"address":"rec5EKMGg6MxZYaMdyBfgwp4d5rB9T1VQH5pJv5LtFJ","metadata":{"name":"price_feed","version":"0.1.0","spec":"0.1.0","description":"Created with Anchor"},"instructions":[],"accounts":[{"name":"PriceUpdateV2","discriminator":[34,241,35,99,157,126,244,205]}],"types":[{"name":"PriceFeedMessage","type":{"kind":"struct","fields":[{"name":"feed_id","docs":[" "],"type":{"array":["u8",32]}},{"name":"price","type":"i64"},{"name":"conf","type":"u64"},{"name":"exponent","type":"i32"},{"name":"publish_time","docs":["The timestamp of this price update in seconds"],"type":"i64"},{"name":"prev_publish_time","docs":[" ",""," "],"type":"i64"},{"name":"ema_price","type":"i64"},{"name":"ema_conf","type":"u64"}]}},{"name":"PriceUpdateV2","type":{"kind":"struct","fields":[{"name":"write_authority","type":"pubkey"},{"name":"verification_level","type":{"defined":{"name":"VerificationLevel"}}},{"name":"price_message","type":{"defined":{"name":"PriceFeedMessage"}}},{"name":"posted_slot","type":"u64"}]}},{"name":"VerificationLevel","type":{"kind":"enum","variants":[{"name":"Partial","fields":[{"name":"num_signatures","type":"u8"}]},{"name":"Full"}]}}]},"address":{"pubkey":"42amVS4KgzR9rA28tkVYqVXjq9Qa8dcZQMbH5EYFX6XC"},"accountType":"PriceUpdateV2","properties":["price_message.price","price_message.publish_time"],"tags":["oracle","price-feed","defi"]},{"id":"pyth-sol-usd-v2","name":"Override SOL/USD Price Feed","description":"Override Pyth SOL/USD price feed with custom price data","protocol":"Pyth","idl":{"address":"rec5EKMGg6MxZYaMdyBfgwp4d5rB9T1VQH5pJv5LtFJ","metadata":{"name":"price_feed","version":"0.1.0","spec":"0.1.0","description":"Created with Anchor"},"instructions":[],"accounts":[{"name":"PriceUpdateV2","discriminator":[34,241,35,99,157,126,244,205]}],"types":[{"name":"PriceFeedMessage","type":{"kind":"struct","fields":[{"name":"feed_id","docs":[" "],"type":{"array":["u8",32]}},{"name":"price","type":"i64"},{"name":"conf","type":"u64"},{"name":"exponent","type":"i32"},{"name":"publish_time","docs":["The timestamp of this price update in seconds"],"type":"i64"},{"name":"prev_publish_time","docs":[" ",""," "],"type":"i64"},{"name":"ema_price","type":"i64"},{"name":"ema_conf","type":"u64"}]}},{"name":"PriceUpdateV2","type":{"kind":"struct","fields":[{"name":"write_authority","type":"pubkey"},{"name":"verification_level","type":{"defined":{"name":"VerificationLevel"}}},{"name":"price_message","type":{"defined":{"name":"PriceFeedMessage"}}},{"name":"posted_slot","type":"u64"}]}},{"name":"VerificationLevel","type":{"kind":"enum","variants":[{"name":"Partial","fields":[{"name":"num_signatures","type":"u8"}]},{"name":"Full"}]}}]},"address":{"pubkey":"7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE"},"accountType":"PriceUpdateV2","properties":["price_message.price","price_message.publish_time"],"tags":["oracle","price-feed","defi"]}] ``` - And the user asks to create a scenario that sets the SOL/USD price to $150.00, the `params` to pass to this method would be: + A valid scenario would look like the user asks to create a scenario that: ```json - { - "id": "scenario-001", - "name": "Set SOL/USD Price to $150", - "description": "A scenario that sets the SOL/USD price feed to $150.00 at slot 5000", - "overrides": [ - { - "id": "override-001", - "templateId": "pyth-sol-usd-v2", - "values": { - "price_message.price": 150000000, // Pyth price is in 10^-6 format - "price_message.publish_time": 1625247600 // Example publish time - }, - "scenarioRelativeSlot": 1, - "label": "Set SOL/USD to $150", - "enabled": true, - "fetchBeforeUse": false, - "account": { - "pubkey": "7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE" - } - } - ], - "tags": ["test", "price-feed"] - } + {"id":"5ec01850-a317-4413-b009-63de4ea10385","name":"New Scenario 2","description":"Add a description...","overrides":[{"id":"pyth_pyth-btc-usd-v2_0","templateId":"pyth_pyth-btc-usd-v2","values":{"price_message.price":12},"scenarioRelativeSlot":1,"label":"Override BTC/USD Price Feed","enabled":true,"fetchBeforeUse":true,"account":{"pubkey":"4cSM2e6rvbGQUFiJbqytoVMi5GgghSMr8LwVrT9VPSPo"}},{"id":"pyth_pyth-eth-btc-v2_0","templateId":"pyth_pyth-eth-btc-v2","values":{"price_message.price":23},"scenarioRelativeSlot":1,"label":"Override ETH/BTC Price Feed","enabled":true,"fetchBeforeUse":true,"account":{"pubkey":"5JwbqPPMNpzE2jVAdobWo6m5gkhsDhRdGBo3FYbSfmaK"}},{"id":"pyth_pyth-eth-btc-v2_1","templateId":"pyth_pyth-eth-btc-v2","values":{"price_message.price":41},"scenarioRelativeSlot":2,"label":"Override ETH/BTC Price Feed","enabled":true,"fetchBeforeUse":true,"account":{"pubkey":"5JwbqPPMNpzE2jVAdobWo6m5gkhsDhRdGBo3FYbSfmaK"}}],"tags":[]} ``` - "#)] - pub fn register_scenario( + pub fn create_scenario( &self, - #[tool(param)] - #[schemars( - description = r#"The parameters to pass to the `surfnet_registerScenario` RPC method. This should contain a `Scenario` object containing a list of `OverrideInstance`s. - Example: - ```json - [ - { - "id": "scenario-001", - "name": "Set SOL/USD Price to $150", - "description": "A scenario that sets the SOL/USD price feed to $150.00 at slot 5000", - "overrides": [ - { - "id": "override-001", - "templateId": "pyth-sol-usd-v2", - "values": { - "price_message.price": 150000000, // Pyth price is in 10^-6 format - "price_message.publish_time": 1625247600 // Example publish time - }, - "scenarioRelativeSlot": 1, - "label": "Set SOL/USD to $150", - "enabled": true, - "fetchBeforeUse": false, - "account": { - "pubkey": "7UVimffxr9ow1uXYxsr4LHAcV58mLzhmwaeKvJ1pjLiE" - } - } - ], - "tags": ["test", "price-feed"] - } - ] - ``` - "# - )] - params: Vec, + #[tool(param)] scenario: Scenario, ) -> Json { - let load_scenarios_endpoint = format!("http://127.0.0.1:{}/v1/scenarios", 18488); - // let params: Vec = params.into_iter().map(Into::into).collect(); - let payload = serde_json::json!(params); + let load_scenarios_endpoint = format!( + "http://127.0.0.1:{}/v1/scenarios", + CHANGE_TO_DEFAULT_STUDIO_PORT_ONCE_SUPERVISOR_MERGED + ); + let payload = serde_json::json!(scenario); // Make the RPC request to the surfnet RPC endpoint let client = reqwest::blocking::Client::new(); @@ -618,13 +539,17 @@ impl Surfpool { ))); } - // Extract the result - let result = rpc_response - .get("result") - .unwrap_or(&serde_json::Value::Null) - .clone(); - - Json(RegisterScenarioResponse::success()) + // Extract the scenario id from the response + let scenario_id = rpc_response + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or(&scenario.id); + + let url = format!( + "http://127.0.0.1:{}/scenarios?id={}&tab=editor", + CHANGE_TO_DEFAULT_STUDIO_PORT_ONCE_SUPERVISOR_MERGED, scenario_id + ); + Json(RegisterScenarioResponse::success(url)) } #[tool( @@ -699,19 +624,19 @@ impl ServerHandler for Surfpool { }) } "str:///override_templates" => { - let pyth_v2_json_value = serde_yaml::from_str::( - PYTH_V2_OVERRIDES_CONTENT, - ) - .expect("Expected Pyth overrides file to be deserializable"); - let joined_value = json!({ - "pyth_v2": pyth_v2_json_value, - }); + let registry = self.template_registry.read().map_err(|_| { + McpError::internal_error("Failed to read template registry", None) + })?; + + let templates = registry.all(); + let templates_json = serde_json::to_string(&templates) + .map_err(|_| McpError::internal_error("Failed to serialize templates", None))?; Ok(ReadResourceResult { contents: vec![ResourceContents::TextResourceContents { uri, mime_type: Some("application/json".to_string()), - text: joined_value.to_string(), + text: templates_json, }], }) } diff --git a/crates/types/Cargo.toml b/crates/types/Cargo.toml index fecfd699..f5decfd3 100644 --- a/crates/types/Cargo.toml +++ b/crates/types/Cargo.toml @@ -18,6 +18,8 @@ blake3 = { workspace = true } chrono = { workspace = true } crossbeam-channel = { workspace = true } once_cell = { workspace = true } +schemars = { workspace = true } +schemars_derive = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } serde_with = { workspace = true, features = ["alloc", "macros"] } diff --git a/crates/types/src/scenarios.rs b/crates/types/src/scenarios.rs index 3e32f531..5a8fb56c 100644 --- a/crates/types/src/scenarios.rs +++ b/crates/types/src/scenarios.rs @@ -11,12 +11,15 @@ use crate::Idl; // ======================================== /// Defines how an account address should be determined -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "camelCase")] +#[doc = "Defines how an account address should be determined"] pub enum AccountAddress { /// A specific public key + #[doc = "A specific public key"] Pubkey(String), /// A Program Derived Address with seeds + #[doc = "A Program Derived Address with seeds"] Pda { program_id: String, seeds: Vec, @@ -24,12 +27,13 @@ pub enum AccountAddress { } /// Seeds used for PDA derivation -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "camelCase")] +#[doc = "Seeds used for PDA derivation"] pub enum PdaSeed { + Pubkey(String), String(String), Bytes(Vec), - Pubkey(String), /// Reference to a property value PropertyRef(String), } @@ -85,26 +89,33 @@ impl OverrideTemplate { } /// A concrete instance of an override template with specific values -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "camelCase")] pub struct OverrideInstance { /// Unique identifier for this instance + #[doc = "Unique identifier for the scenario"] pub id: String, /// Reference to the template being used + #[doc = "Reference to the template being used"] pub template_id: String, /// Values for the template properties (flat key-value map with dot notation, e.g., "price_message.price_value") + #[doc = "Values for the template properties (flat key-value map with dot notation, e.g., 'price_message.price_value')"] pub values: HashMap, /// Relative slot when this override should be applied (relative to scenario registration slot) + #[doc = "Relative slot when this override should be applied (relative to scenario registration slot)"] pub scenario_relative_slot: Slot, /// Optional label for this instance + #[doc = "Optional label for this instance"] pub label: Option, /// Whether this override is enabled + #[doc = "Whether this override is enabled"] pub enabled: bool, /// Whether to fetch fresh account data just before transaction execution - /// Useful for time-sensitive data like price feeds, oracle updates, and dynamic balances + #[doc = "Whether to fetch fresh account data just before transaction execution"] #[serde(default)] pub fetch_before_use: bool, /// Account address to override + #[doc = "Account address to override"] pub account: AccountAddress, } @@ -134,18 +145,23 @@ impl OverrideInstance { } /// A scenario containing a timeline of overrides -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, schemars::JsonSchema)] #[serde(rename_all = "camelCase")] pub struct Scenario { /// Unique identifier for the scenario + #[doc = "Unique identifier for the scenario"] pub id: String, /// Human-readable name + #[doc = "Human-readable name"] pub name: String, /// Description of this scenario + #[doc = "Description of this scenario"] pub description: String, /// List of override instances in this scenario + #[doc = "List of override instances in this scenario"] pub overrides: Vec, /// Tags for categorization + #[doc = "Tags for categorization"] pub tags: Vec, } From 9200da6e8d292c1d630b51cd98f5b8c418300fbf Mon Sep 17 00:00:00 2001 From: Ludo Galabru Date: Thu, 13 Nov 2025 23:11:09 -0500 Subject: [PATCH 5/5] feat: snapshot scopes --- crates/core/src/rpc/surfnet_cheatcodes.rs | 141 +++++++++++++++++++++- crates/core/src/surfnet/svm.rs | 51 ++++++-- crates/types/src/types.rs | 53 ++++++++ 3 files changed, 235 insertions(+), 10 deletions(-) diff --git a/crates/core/src/rpc/surfnet_cheatcodes.rs b/crates/core/src/rpc/surfnet_cheatcodes.rs index 562f02c1..16eca302 100644 --- a/crates/core/src/rpc/surfnet_cheatcodes.rs +++ b/crates/core/src/rpc/surfnet_cheatcodes.rs @@ -792,6 +792,9 @@ pub trait SurfnetCheatcodes { /// - `includeProgramAccounts`: A boolean indicating whether to include program accounts. /// - `includeAccounts`: A list of specific account public keys to include. /// - `excludeAccounts`: A list of specific account public keys to exclude. + /// - `scope`: An optional scope to limit the accounts included in the snapshot. Options include: + /// - `network`: Includes all accounts in the network. + /// - `preTransaction`: Only includes accounts touched by the given transaction. /// /// /// ## Returns @@ -1834,7 +1837,8 @@ mod tests { use spl_token_2022_interface::instruction::{initialize_mint2, mint_to, transfer_checked}; use spl_token_interface::state::Mint; use surfpool_types::{ - ExportSnapshotFilter, RpcProfileDepth, UiAccountChange, UiAccountProfileState, + ExportSnapshotFilter, ExportSnapshotScope, RpcProfileDepth, UiAccountChange, + UiAccountProfileState, }; use super::*; @@ -2755,6 +2759,7 @@ mod tests { Some(ExportSnapshotConfig { include_parsed_accounts: Some(true), filter: None, + scope: ExportSnapshotScope::Network, }), ) .expect("Failed to export snapshot") @@ -2798,6 +2803,140 @@ mod tests { ); } + #[test] + fn test_export_snapshot_pre_transaction() { + use std::collections::HashMap; + + use solana_signature::Signature; + use surfpool_types::{ProfileResult, types::KeyedProfileResult}; + + let client = TestSetup::new(SurfnetCheatcodesRpc); + + // Create several accounts in the network + let account1_pubkey = Pubkey::new_unique(); + let account1 = Account { + lamports: 1_000_000, + data: vec![1, 2, 3, 4], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + set_account(&client, &account1_pubkey, &account1); + + let account2_pubkey = Pubkey::new_unique(); + let account2 = Account { + lamports: 2_000_000, + data: vec![5, 6, 7, 8], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + set_account(&client, &account2_pubkey, &account2); + + let account3_pubkey = Pubkey::new_unique(); + let account3 = Account { + lamports: 3_000_000, + data: vec![9, 10, 11, 12], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + set_account(&client, &account3_pubkey, &account3); + + // Create a mock transaction profile that only touches account1 and account2 + let signature = Signature::new_unique(); + let mut pre_execution_capture = BTreeMap::new(); + pre_execution_capture.insert(account1_pubkey, Some(account1.clone())); + pre_execution_capture.insert(account2_pubkey, Some(account2.clone())); + + let mut post_execution_capture = BTreeMap::new(); + let mut modified_account1 = account1.clone(); + modified_account1.lamports = 500_000; + post_execution_capture.insert(account1_pubkey, Some(modified_account1.clone())); + let mut modified_account2 = account2.clone(); + modified_account2.lamports = 2_500_000; + post_execution_capture.insert(account2_pubkey, Some(modified_account2.clone())); + + let profile = ProfileResult { + pre_execution_capture, + post_execution_capture, + compute_units_consumed: 1000, + log_messages: None, + error_message: None, + }; + + let keyed_profile = KeyedProfileResult::new( + 1, + UuidOrSignature::Signature(signature), + None, + profile, + HashMap::new(), + ); + + // Insert the profile into executed_transaction_profiles + client.context.svm_locker.with_svm_writer(|svm| { + svm.executed_transaction_profiles + .insert(signature, keyed_profile); + }); + + // Export snapshot with PreTransaction scope + let snapshot = client + .rpc + .export_snapshot( + Some(client.context.clone()), + Some(ExportSnapshotConfig { + include_parsed_accounts: Some(false), + filter: None, + scope: ExportSnapshotScope::PreTransaction(signature.to_string()), + }), + ) + .expect("Failed to export snapshot") + .value; + + // Verify that only account1 and account2 are in the snapshot + assert!( + snapshot.contains_key(&account1_pubkey.to_string()), + "Snapshot should contain account1 (touched by transaction)" + ); + assert!( + snapshot.contains_key(&account2_pubkey.to_string()), + "Snapshot should contain account2 (touched by transaction)" + ); + assert!( + !snapshot.contains_key(&account3_pubkey.to_string()), + "Snapshot should NOT contain account3 (not touched by transaction)" + ); + + // Verify the accounts have the PRE-EXECUTION state (original values, not modified) + verify_snapshot_account(&snapshot, &account1_pubkey, &account1); + verify_snapshot_account(&snapshot, &account2_pubkey, &account2); + + // Double-check that we're NOT getting the post-execution values + let snapshot_account1 = snapshot + .get(&account1_pubkey.to_string()) + .expect("Account1 should be in snapshot"); + assert_eq!( + snapshot_account1.lamports, 1_000_000, + "Account1 should have pre-execution lamports (1M), not post-execution (500K)" + ); + + let snapshot_account2 = snapshot + .get(&account2_pubkey.to_string()) + .expect("Account2 should be in snapshot"); + assert_eq!( + snapshot_account2.lamports, 2_000_000, + "Account2 should have pre-execution lamports (2M), not post-execution (2.5M)" + ); + + // Verify account count + // Note: The snapshot may contain more accounts if system accounts are included + // but we verify that at least our touched accounts are there and untouched ones are not + println!( + "Snapshot contains {} accounts (expected at least 2)", + snapshot.len() + ); + } + #[test] fn test_export_snapshot_filtering() { let system_account_pubkey = Pubkey::new_unique(); diff --git a/crates/core/src/surfnet/svm.rs b/crates/core/src/surfnet/svm.rs index acbb734e..c9c87eb9 100644 --- a/crates/core/src/surfnet/svm.rs +++ b/crates/core/src/surfnet/svm.rs @@ -52,10 +52,10 @@ use spl_token_2022_interface::extension::{ }; use surfpool_types::{ AccountChange, AccountProfileState, AccountSnapshot, DEFAULT_PROFILING_MAP_CAPACITY, - DEFAULT_SLOT_TIME_MS, ExportSnapshotConfig, FifoMap, Idl, OverrideInstance, ProfileResult, - RpcProfileDepth, RpcProfileResultConfig, RunbookExecutionStatusReport, SimnetEvent, - TransactionConfirmationStatus, TransactionStatusEvent, UiAccountChange, UiAccountProfileState, - UiProfileResult, VersionedIdl, + DEFAULT_SLOT_TIME_MS, ExportSnapshotConfig, ExportSnapshotScope, FifoMap, Idl, + OverrideInstance, ProfileResult, RpcProfileDepth, RpcProfileResultConfig, + RunbookExecutionStatusReport, SimnetEvent, TransactionConfirmationStatus, + TransactionStatusEvent, UiAccountChange, UiAccountProfileState, UiProfileResult, VersionedIdl, types::{ ComputeUnitsEstimationResult, KeyedProfileResult, UiKeyedProfileResult, UuidOrSignature, }, @@ -2323,16 +2323,17 @@ impl SurfnetSvm { || pubkey == &solana_sdk_ids::bpf_loader_deprecated::id() || pubkey == &solana_sdk_ids::bpf_loader_upgradeable::id() } - for (pubkey, account_shared_data) in self.iter_accounts() { + + // Helper function to process an account and add it to fixtures + let mut process_account = |pubkey: &Pubkey, account: &Account| { let is_include_account = include_accounts.iter().any(|k| k.eq(&pubkey.to_string())); let is_exclude_account = exclude_accounts.iter().any(|k| k.eq(&pubkey.to_string())); - let is_program_account = is_program_account(account_shared_data.owner()); + let is_program_account = is_program_account(&account.owner); if is_exclude_account || ((is_program_account && !include_program_accounts) && !is_include_account) { - continue; + return; } - let account = Account::from(account_shared_data.clone()); // For token accounts, we need to provide the mint additional data let additional_data = if account.owner == spl_token_interface::id() @@ -2350,7 +2351,7 @@ impl SurfnetSvm { }; let ui_account = - self.encode_ui_account(pubkey, &account, encoding, additional_data, None); + self.encode_ui_account(pubkey, account, encoding, additional_data, None); let (base64, parsed_data) = match ui_account.data { UiAccountData::Json(parsed_account) => { @@ -2370,7 +2371,39 @@ impl SurfnetSvm { ); fixtures.insert(pubkey.to_string(), account_snapshot); + }; + + match &config.scope { + ExportSnapshotScope::Network => { + // Export all network accounts (current behavior) + for (pubkey, account_shared_data) in self.iter_accounts() { + let account = Account::from(account_shared_data.clone()); + process_account(&pubkey, &account); + } + } + ExportSnapshotScope::PreTransaction(signature_str) => { + // Export accounts from a specific transaction's pre-execution state + if let Ok(signature) = Signature::from_str(signature_str) { + if let Some(profile) = self.executed_transaction_profiles.get(&signature) { + // Collect accounts from pre-execution capture only + // This gives us the account state BEFORE the transaction executed + for (pubkey, account_opt) in + &profile.transaction_profile.pre_execution_capture + { + if let Some(account) = account_opt { + process_account(pubkey, account); + } + } + + // Also collect readonly account states (these don't change) + for (pubkey, account) in &profile.readonly_account_states { + process_account(pubkey, account); + } + } + } + } } + fixtures } diff --git a/crates/types/src/types.rs b/crates/types/src/types.rs index 16473bab..91baa4d3 100644 --- a/crates/types/src/types.rs +++ b/crates/types/src/types.rs @@ -1002,6 +1002,15 @@ impl AccountSnapshot { pub struct ExportSnapshotConfig { pub include_parsed_accounts: Option, pub filter: Option, + pub scope: ExportSnapshotScope, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum ExportSnapshotScope { + #[default] + Network, + PreTransaction(String), } #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] @@ -1378,4 +1387,48 @@ mod tests { println!("Profiling map: {:?}", get); assert_eq!(get, vec![("b", 4), ("c", 3), ("d", 4), ("e", 5)]); } + + #[test] + fn test_export_snapshot_scope_serialization() { + // Test Network variant + let network_config = ExportSnapshotConfig { + include_parsed_accounts: None, + filter: None, + scope: ExportSnapshotScope::Network, + }; + let network_json = serde_json::to_value(&network_config).unwrap(); + println!( + "Network config: {}", + serde_json::to_string_pretty(&network_json).unwrap() + ); + assert_eq!(network_json["scope"], json!("network")); + + // Test PreTransaction variant + let pre_tx_config = ExportSnapshotConfig { + include_parsed_accounts: None, + filter: None, + scope: ExportSnapshotScope::PreTransaction("5signature123".to_string()), + }; + let pre_tx_json = serde_json::to_value(&pre_tx_config).unwrap(); + println!( + "PreTransaction config: {}", + serde_json::to_string_pretty(&pre_tx_json).unwrap() + ); + assert_eq!( + pre_tx_json["scope"], + json!({"preTransaction": "5signature123"}) + ); + + // Test deserialization + let deserialized_network: ExportSnapshotConfig = + serde_json::from_value(network_json).unwrap(); + assert_eq!(deserialized_network.scope, ExportSnapshotScope::Network); + + let deserialized_pre_tx: ExportSnapshotConfig = + serde_json::from_value(pre_tx_json).unwrap(); + assert_eq!( + deserialized_pre_tx.scope, + ExportSnapshotScope::PreTransaction("5signature123".to_string()) + ); + } }