Strategy split ws

src/engine/engine/mod.rs

@@ -7,6 +7,7 @@ use crate::{
     store::{accounts::AccountList, config::Config},
     terminal::TerminalServer,
 };
+use hypersdk::hypercore::ws::ConnectionStream;
 use pulse_sdk::prelude::*;
 use std::{collections::HashMap, sync::Arc};
 use tokio::{sync::Mutex, task::JoinHandle};
@@ -19,8 +20,12 @@ pub struct WatchList {
 
 #[derive(Clone)]
 pub struct Engine {
+    // engine
     pub terminal_server: Arc<TerminalServer>,
     pub strategy_engine: Arc<StrategyEngine>,
+    pub ws_stream: Arc<Mutex<ConnectionStream>>,
+
+    // data
     pub config: Arc<Mutex<Config>>,
     pub accounts: Arc<Mutex<AccountList>>,
     pub watch_list: Arc<Mutex<WatchList>>,
@@ -31,16 +36,19 @@ impl Engine {
     pub async fn new() -> tokio::io::Result<Arc<Self>> {
         let config = Config::new().await?;
 
-        let strategy = StrategyEngine::new(&config.strategy).await?;
+        let (ws_handle, ws_stream) = hypersdk::hypercore::mainnet_ws().split();
+
+        let strategy = StrategyEngine::new(&config.strategy, ws_handle).await?;
 
         let accounts = Arc::new(Mutex::new(AccountList::new().await?));
         let config = Arc::new(Mutex::new(config));
 
         Ok(Arc::new_cyclic(|engine| Self {
-            terminal_server: TerminalServer::new(engine.clone()),
-            strategy_engine: strategy.initialize(engine.clone()),
             config,
             accounts,
+            ws_stream: Arc::new(Mutex::new(ws_stream)),
+            terminal_server: TerminalServer::new(engine.clone()),
+            strategy_engine: strategy.initialize(engine.clone()),
             watch_list: Arc::new(Mutex::new(WatchList {
                 name_to_index: HashMap::new(),
                 items: Vec::new(),

src/engine/engine/strategy.rs

@@ -1,5 +1,5 @@
 use anyhow::Context;
-use hypersdk::hypercore::{self, CandleInterval, Subscription, WebSocket};
+use hypersdk::hypercore::{self, CandleInterval, Subscription, ws::ConnectionHandle};
 use pulse_sdk::prelude::*;
 use std::{
     collections::HashSet,
@@ -19,12 +19,12 @@ use crate::{
 pub struct StrategyEngine {
     pub engine: Weak<Engine>,
     pub strategy: Mutex<StrategyChild>,
-    pub ws: WebSocket,
+    pub ws_handle: ConnectionHandle,
     pub subscriptions: Mutex<HashSet<Subscription>>,
 }
 
 impl StrategyEngine {
-    pub async fn new(strategy_id: &str) -> tokio::io::Result<Self> {
+    pub async fn new(strategy_id: &str, ws_handle: ConnectionHandle) -> tokio::io::Result<Self> {
         let strategy = pulse_strategy(strategy_id)?;
 
         let (strategy, strategy_manifest) = get_manifest(
@@ -42,7 +42,7 @@ impl StrategyEngine {
         Ok(Self {
             strategy: Mutex::new(StrategyChild::new(strategy, strategy_manifest)),
             engine: Weak::new(),
-            ws: hypercore::mainnet_ws(),
+            ws_handle,
             subscriptions: Mutex::new(HashSet::new()),
         })
     }
@@ -142,18 +142,18 @@ impl StrategyEngine {
                 }
 
                 Some(StrategyMessage::Subscribe(subscription)) => {
-                    self.ws.subscribe(subscription.clone());
+                    self.ws_handle.subscribe(subscription.clone());
                     self.subscriptions.lock().await.insert(subscription);
                 }
 
                 Some(StrategyMessage::Unsubscribe(subscription)) => {
                     self.subscriptions.lock().await.remove(&subscription);
-                    self.ws.unsubscribe(subscription);
+                    self.ws_handle.unsubscribe(subscription);
                 }
 
                 Some(StrategyMessage::UnsubscribeAll) => {
                     for sub in self.subscriptions.lock().await.drain() {
-                        self.ws.unsubscribe(sub);
+                        self.ws_handle.unsubscribe(sub);
                     }
                 }