From 1dc1402cd21f2f3ca56c439e795bd954a360fc85 Mon Sep 17 00:00:00 2001 From: Klesti Selimaj Date: Thu, 6 Aug 2026 21:46:56 +0200 Subject: [PATCH] Strategy structure --- src/bot.rs | 65 +++++++++++++------------------------------ src/main.rs | 1 + src/strategy/mod.rs | 15 ++++++++++ src/strategy/strat.rs | 47 +++++++++++++++++++++++++++++++ 4 files changed, 83 insertions(+), 45 deletions(-) create mode 100644 src/strategy/mod.rs create mode 100644 src/strategy/strat.rs diff --git a/src/bot.rs b/src/bot.rs index b8ccd13..7b74216 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use anyhow::Context; use futures_util::{SinkExt, StreamExt}; @@ -9,52 +9,17 @@ use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use crate::{ account::AccountManager, executor::Executor, - types::{Mode, NewToken, Token, Trade}, + strategy::{Strategy, strat::Strat}, }; pub struct Bot { pub ws: Mutex>>, pub accounts: Mutex, pub executor: Mutex>, - pub tokens: Mutex>, + pub strategy: Mutex>, } impl Bot { - pub async fn on_new_coin(&self, token: NewToken) -> anyhow::Result<()> { - self.tokens.lock().await.insert( - token.mint.clone(), - Token { - mode: Mode::Observing, - execute_next: false, - }, - ); - - self.subscribe(&token.mint).await?; - - Ok(()) - } - - pub async fn on_trade(&self, trade: Trade) -> anyhow::Result<()> { - log::info!("Trade: {trade:?}"); - - let mut tokens = self.tokens.lock().await; - - let Some(token) = tokens.get_mut(&trade.mint) else { - log::error!("Token not found on trade: {:?}", trade.mint); - return Ok(()); - }; - - match &token.mode { - Mode::Observing => {} - - Mode::WaitingForEntry => {} - - Mode::WaitingForExit => {} - } - - Ok(()) - } - pub async fn subscribe(&self, mint: &str) -> anyhow::Result<()> { self.ws .lock() @@ -91,7 +56,7 @@ impl Bot { } impl Bot { - pub async fn new() -> anyhow::Result { + pub async fn new() -> anyhow::Result> { let accounts = AccountManager::get().await?; let account = accounts @@ -100,12 +65,14 @@ impl Bot { .context("Failed to get account")? .clone(); - Ok(Self { + Ok(Arc::new(Self { ws: Mutex::new(connect_async("wss://pumpdev.io/ws").await?.0), executor: Mutex::new(Box::new(account.executor())), accounts: Mutex::new(accounts), - tokens: Mutex::new(HashMap::new()), - }) + strategy: Mutex::new(Box::new(Strat { + tokens: HashMap::new(), + })), + })) } pub async fn refresh_account(&self) -> anyhow::Result<()> { @@ -134,7 +101,7 @@ impl Bot { Ok(()) } - pub async fn start(&self) -> anyhow::Result<()> { + pub async fn start(self: &Arc) -> anyhow::Result<()> { self.initialize_websocket_subscribe().await?; let mut ws = self.ws.lock().await; @@ -144,13 +111,21 @@ impl Bot { match serde_json::from_str::(&text) { Ok(crate::types::PumpDevEvent::Create(token)) => { drop(ws); - self.on_new_coin(token).await?; + self.strategy + .lock() + .await + .on_new_coin(self.clone(), token) + .await?; ws = self.ws.lock().await; } Ok(crate::types::PumpDevEvent::Trade(trade)) => { drop(ws); - self.on_trade(trade).await?; + self.strategy + .lock() + .await + .on_trade(self.clone(), trade) + .await?; ws = self.ws.lock().await; } diff --git a/src/main.rs b/src/main.rs index 73ba992..c929842 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ pub mod account; pub mod bot; pub mod executor; +pub mod strategy; pub mod types; use crate::bot::Bot; diff --git a/src/strategy/mod.rs b/src/strategy/mod.rs new file mode 100644 index 0000000..3e47490 --- /dev/null +++ b/src/strategy/mod.rs @@ -0,0 +1,15 @@ +pub mod strat; + +use std::sync::Arc; + +use crate::{ + bot::Bot, + types::{NewToken, Trade}, +}; + +#[async_trait::async_trait] +pub trait Strategy: Send + Sync { + async fn on_new_coin(&mut self, bot: Arc, token: NewToken) -> anyhow::Result<()>; + + async fn on_trade(&mut self, bot: Arc, trade: Trade) -> anyhow::Result<()>; +} diff --git a/src/strategy/strat.rs b/src/strategy/strat.rs new file mode 100644 index 0000000..5a7ed87 --- /dev/null +++ b/src/strategy/strat.rs @@ -0,0 +1,47 @@ +use std::{collections::HashMap, sync::Arc}; + +use crate::{ + bot::Bot, + strategy::Strategy, + types::{Mode, NewToken, Token, Trade}, +}; + +pub struct Strat { + pub tokens: HashMap, +} + +#[async_trait::async_trait] +impl Strategy for Strat { + async fn on_new_coin(&mut self, bot: Arc, token: NewToken) -> anyhow::Result<()> { + self.tokens.insert( + token.mint.clone(), + Token { + mode: Mode::Observing, + execute_next: false, + }, + ); + + bot.subscribe(&token.mint).await?; + + Ok(()) + } + + async fn on_trade(&mut self, bot: Arc, trade: Trade) -> anyhow::Result<()> { + log::info!("Trade: {trade:?}"); + + let Some(token) = self.tokens.get_mut(&trade.mint) else { + log::error!("Token not found on trade: {:?}", trade.mint); + return Ok(()); + }; + + match &token.mode { + Mode::Observing => {} + + Mode::WaitingForEntry => {} + + Mode::WaitingForExit => {} + } + + Ok(()) + } +}