Agentic refactoring and full MCP support #1502

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

gustavo-grieco wants to merge 40 commits into master from dev-agents

flake.nix

-Original file line number
+Diff line change
@@ Expand Up / @@ -69,8 +69,15 @@ @@
                 pkgs.haskell.lib.compose.dontCheck
               ]);
+            mcp-server = pkgs: pkgs.haskellPackages.callCabal2nix "mcp-server" (pkgs.fetchFromGitHub {
+              owner = "gustavo-grieco";
+              repo = "haskell-mcp-server";
+              rev = "9fd60af428b96ae4bc63a133b3960ed934494189";
+              sha256 = "sha256-lh65Gy8a43xbDDFPONOJ2UBUS1xWOW2UUx3wYFTG8Xg=";
+            }) {};
             echidna = pkgs: with pkgs; lib.pipe
-              (haskellPackages.callCabal2nix "echidna" ./. { hevm = hevm pkgs; })
+              (haskellPackages.callCabal2nix "echidna" ./. { hevm = hevm pkgs; mcp-server = mcp-server pkgs; })
               ([
                 # FIXME: figure out solc situation, it conflicts with the one from
                 # solc-select that is installed with slither, disable tests in the meantime
@@ Expand Down @@

lib/Echidna.hs

-Original file line number
+Diff line change
@@ -1,6 +1,7 @@
     module Echidna where
     import Control.Concurrent (newChan)
+    import Control.Concurrent.STM (newBroadcastTChanIO)
     import Control.Monad.Catch (MonadThrow(..))
     import Control.Monad.IO.Class (liftIO)
     import Data.IORef (newIORef)
@@ Expand Down Expand Up @@
                     <> staticAddresses solConf
                     <> deployedAddresses
         deployedSolcContracts = nub $ mapMaybe (findSrcForReal env.dapp) $ Map.elems vm.env.contracts
-        nonViewPureSigs = concatMap (mapMaybe (\ (Method {name, inputs, mutability}) ->
+        nonViewPureSigs = concatMap (mapMaybe (\ (Method {name, inputs, mutability}) ->
           case mutability of
             View -> Nothing
             Pure -> Nothing
@@ Expand Down Expand Up / @@ -123,6 +124,7 @@ mkEnv cfg buildOutput tests world slitherInfo = do @@
       codehashMap <- newIORef mempty
       chainId <- Onchain.fetchChainIdFrom cfg.rpcUrl
       eventQueue <- newChan
+      bus <- newBroadcastTChanIO
       coverageRefInit <- newIORef mempty
       coverageRefRuntime <- newIORef mempty
       corpusRef <- newIORef mempty
@@ Expand All / @@ -131,7 +133,8 @@ mkEnv cfg buildOutput tests world slitherInfo = do @@
       contractNameCache <- newIORef mempty
       -- TODO put in real path
       let dapp = dappInfo "/" buildOutput
-      pure $ Env { cfg, dapp, codehashMap, fetchSession, contractNameCache
-                 , chainId, eventQueue, coverageRefInit, coverageRefRuntime, corpusRef, testRefs, world
+          sourceCache = buildOutput.sources
+      pure $ Env { cfg, dapp, sourceCache, codehashMap, fetchSession, contractNameCache
+                 , chainId, eventQueue, bus, coverageRefInit, coverageRefRuntime, corpusRef, testRefs, world
                  , slitherInfo
                  }

lib/Echidna/Agent/Fuzzer.hs

-Original file line number
+Diff line change
@@ -0,0 +1,325 @@
+    {-# LANGUAGE GADTs #-}
+    {-# LANGUAGE DataKinds #-}
+    {-# LANGUAGE NamedFieldPuns #-}
+    {-# LANGUAGE OverloadedRecordDot #-}
+    {-# LANGUAGE MultiWayIf #-}
+    module Echidna.Agent.Fuzzer where
+    import Control.Concurrent.STM (atomically, tryReadTChan, dupTChan, putTMVar)
+    import Control.Monad (replicateM, void, forM_, when)
+    import Control.Monad.Reader (runReaderT, liftIO, asks, MonadReader, ask)
+    import Control.Monad.State.Strict (runStateT, get, gets, modify', MonadState)
+    import Control.Monad.Random.Strict (evalRandT, MonadRandom, RandT, getRandom, getRandomR)
+    import Control.Monad.Catch (MonadThrow(..))
+    import Control.Monad.Trans (lift)
+    import Control.Monad.IO.Class (MonadIO)
+    import System.Random (mkStdGen)
+    import Data.IORef (IORef, writeIORef, readIORef, atomicModifyIORef')
+    import Data.Map (Map)
+    import qualified Data.Map as Map
+    import qualified Data.Set as Set
+    import Data.Text (Text)
+    import System.Directory (getCurrentDirectory)
+    import Echidna.Output.Source (saveLcovHook)
+    import EVM.Dapp (DappInfo(..))
+    import EVM.Types (VM(..), VMType(Concrete), Expr(..), EType(..), Contract)
+    import qualified EVM.Types as EVM
+    import EVM.ABI (AbiValue)
+    import Echidna.ABI (GenDict(..))
+    import Echidna.Execution (replayCorpus, callseq, updateTests)
+    import Echidna.Mutator.Corpus (getCorpusMutation, seqMutatorsStateless, seqMutatorsStateful, fromConsts)
+    import Echidna.Shrink (shrinkTest)
+    import Echidna.Transaction (genTx, genTxFromPrototype)
+    import Echidna.Types.Random (rElem)
+    import qualified Data.List.NonEmpty as NE
+    import Echidna.Types.Agent
+    import Echidna.Types.Campaign (WorkerState(..), CampaignConf(..))
+    import Echidna.Types.Config (Env(..), EConfig(..))
+    import Echidna.Types.InterWorker (AgentId(..), Bus, WrappedMessage(..), Message(..), FuzzerCmd(..))
+    import Echidna.Types.Test (EchidnaTest(..), TestState(..), TestType(..), isOpen, isOptimizationTest)
+    import Echidna.Types.Tx (Tx)
+    import Echidna.Types.Worker (WorkerEvent(..), WorkerType(..), CampaignEvent(..), WorkerStopReason(..))
+    import qualified Echidna.Types.Worker as Worker
+    import Echidna.Worker (pushCampaignEvent)
+    instance (MonadThrow m) => MonadThrow (RandT g m) where
+      throwM = lift . throwM
+    data FuzzerAgent = FuzzerAgent
+      { fuzzerId :: Int
+      , initialVm :: VM Concrete
+      , initialDict :: GenDict
+      , initialCorpus :: [(FilePath, [Tx])]
+      , testLimit :: Int
+      , stateRef :: IORef WorkerState
+      }
+    instance Show FuzzerAgent where
+      show agent = "FuzzerAgent { fuzzerId = " ++ show agent.fuzzerId ++ " }"
+    instance Agent FuzzerAgent where
+      getAgentId agent = FuzzerId agent.fuzzerId
+      runAgent agent bus env = do
+        let workerId = agent.fuzzerId
+            vm = agent.initialVm
+            dict = agent.initialDict
+            corpus = agent.initialCorpus
+            limit = agent.testLimit
+            ref = agent.stateRef
+            effectiveSeed = dict.defSeed + workerId
+            effectiveGenDict = dict { defSeed = effectiveSeed }
+            initialState = WorkerState
+              { workerId
+              , genDict = effectiveGenDict
+              , newCoverage = False
+              , ncallseqs = 0
+              , ncalls = 0
+              , totalGas = 0
+              , runningThreads = []
+              , prioritizedSequences = []
+              }
+        -- Callback to update the IORef with the current state
+        let callback = get >>= liftIO . writeIORef ref
+        (reason, _) <- flip evalRandT (mkStdGen effectiveSeed) $
+               flip runReaderT env $
+                 flip runStateT initialState $ do
+                   liftIO $ pushCampaignEvent env (WorkerEvent workerId FuzzWorker (Worker.Log ("Starting FuzzerAgent " ++ show workerId)))
+                   callback
+                   void $ replayCorpus vm corpus
+                   workerChan <- liftIO $ atomically $ dupTChan bus
+                   fuzzerLoop callback vm limit workerChan
+        liftIO $ pushCampaignEvent env (WorkerEvent workerId FuzzWorker (WorkerStopped reason))
+        return ()
+    fuzzerLoop
+      :: (MonadIO m, MonadThrow m, MonadRandom m, MonadReader Env m, MonadState WorkerState m)
+      => m () -- ^ Callback
+      -> VM Concrete
+      -> Int -- ^ Test limit
+      -> Bus
+      -> m WorkerStopReason
+    fuzzerLoop callback vm testLimit bus = do
+      -- Check for messages
+      -- TODO: Properly handle messages. For now we just check if we should stop?
+      -- But runAgent doesn't return until done.
+      -- We can peek the bus. But standard fuzzer might be busy.
+      -- Maybe check bus every N iterations?
+      run
+      where
+      run = do
+        checkMessages
+        testRefs <- asks (.testRefs)
+        tests <- liftIO $ traverse readIORef testRefs
+        CampaignConf{stopOnFail, shrinkLimit} <- asks (.cfg.campaignConf)
+        ncalls <- gets (.ncalls)
+        workerId <- gets (.workerId)
+        let
+          shrinkable test =
+            case test.state of
+              -- we shrink only tests which were solved on this
+              -- worker
+              Large n | test.workerId == Just workerId ->
+                n < shrinkLimit
+              _       -> False
+          final test =
+            case test.state of
+              Solved   -> True
+              Failed _ -> True
+              _        -> False
+          closeOptimizationTest test =
+            case test.testType of
+              OptimizationTest _ _ ->
+                test { state = Large 0
+                     , workerId = Just workerId
+                     }
+              _ -> test
+        if | stopOnFail && any final tests ->
+             callback >> pure FastFailed
+           -- we shrink first before going back to fuzzing
+           | any shrinkable tests ->
+             shrink >> callback >> run
+           -- no shrinking work, fuzz
+           | (null tests || any isOpen tests) && ncalls < testLimit ->
+             fuzz >> callback >> run
+           -- NOTE: this is a hack which forces shrinking of optimization tests
+           -- after test limit is reached
+           | ncalls >= testLimit && any (\t -> isOpen t && isOptimizationTest t) tests -> do
+             liftIO $ forM_ testRefs $ \testRef ->
+                atomicModifyIORef' testRef (\test -> (closeOptimizationTest test, ()))
+             callback >> run
+           -- no more work to do, means we reached the test limit, exit
+           | otherwise ->
+             callback >> pure TestLimitReached
+      fuzz = randseq vm.env.contracts >>= fmap fst . (\txs -> callseq vm txs False)
+      shrink = do
+        wid <- gets (.workerId)
+        updateTests $ \test -> do
+          if test.workerId == Just wid then
+            shrinkTest vm test
+          else
+            pure Nothing
+      checkMessages = do
+         -- Non-blocking read
+         msg <- liftIO $ atomically $ tryReadTChan bus
+         case msg of
+           Just (WrappedMessage _ (ToFuzzer tid (SolutionFound _))) -> do
+              workerId <- gets (.workerId)
+              when (tid == workerId) $ do
+                 -- Received help!
+                 pure ()
+           Just (WrappedMessage _ (ToFuzzer tid DumpLcov)) -> do
+              workerId <- gets (.workerId)
+              when (tid == workerId) $ do
+                env <- ask
+                liftIO $ do
+                   let contracts = Map.elems env.dapp.solcByName
+                   dir <- maybe getCurrentDirectory pure env.cfg.campaignConf.corpusDir
+                   void $ saveLcovHook env dir env.sourceCache contracts
+                   putStrLn $ "Fuzzer " ++ show workerId ++ ": dumped LCOV coverage."
+                pure ()
+           Just (WrappedMessage _ (ToFuzzer tid (FuzzSequence txs prob))) -> do
+              workerId <- gets (.workerId)
+              when (tid == workerId) $ do
+                 modify' $ \s -> s { prioritizedSequences = (prob, txs) : s.prioritizedSequences }
+                 pure ()
+           Just (WrappedMessage _ (ToFuzzer tid ClearPrioritization)) -> do
+              workerId <- gets (.workerId)
+              when (tid == workerId) $ do
+                 modify' $ \s -> s { prioritizedSequences = [] }
+                 pure ()
+           Just (WrappedMessage _ (ToFuzzer tid (ExecuteSequence txs replyVar))) -> do
+              workerId <- gets (.workerId)
+              when (tid == workerId) $ do
+                 (_, newCov) <- callseq vm txs False
+                 liftIO $ case replyVar of
+                    Just var -> atomically $ putTMVar var newCov
+                    Nothing -> pure ()
+                 pure ()
+           _ -> pure ()
+    -- | Generate a new sequences of transactions, either using the corpus or with
+    -- randomly created transactions
+    randseq
+      :: (MonadRandom m, MonadReader Env m, MonadState WorkerState m, MonadIO m)
+      => Map (Expr 'EAddr) Contract
+      -> m [Tx]
+    randseq deployedContracts = do
+      -- 1. Check for prioritized sequences injected via tools
+      prioritized <- gets (.prioritizedSequences)
+      mbSeq <- if null prioritized
+               then pure Nothing
+               else do
+                 -- Select a prioritized sequence based on probability
+                 (prob, seqPrototype) <- rElem (NE.fromList prioritized)
+                 useIt <- (<= prob) <$> getRandom
+                 pure $ if useIt then Just seqPrototype else Nothing
+      case mbSeq of
+        Just seqPrototype -> genPrioritizedSeq deployedContracts seqPrototype
+        Nothing -> genStandardSeq deployedContracts
+    -- | Generate a sequence of transactions based on a prioritized prototype
+    genPrioritizedSeq
+      :: (MonadRandom m, MonadReader Env m, MonadState WorkerState m, MonadIO m)
+      => Map (Expr 'EAddr) Contract
+      -> [(Text, [Maybe AbiValue])]
+      -> m [Tx]
+    genPrioritizedSeq deployedContracts seqPrototype = do
+           env <- ask
+           let world = env.world
+               seqLen = env.cfg.campaignConf.seqLen
+           -- 2. If a prioritized sequence is selected:
+           -- Expand the prototype into concrete transactions
+           let expandPrototype [] = return []
+               expandPrototype [p] = do
+                   tx <- genTxFromPrototype world deployedContracts p
+                   return [tx]
+               expandPrototype (p:ps) = do
+                   tx <- genTxFromPrototype world deployedContracts p
+                   -- Insert random transactions between prototype transactions to increase fuzzing diversity
+                   n <- getRandomR (0, 3)
+                   rndTxs <- replicateM n (genTx world deployedContracts)
+                   rest <- expandPrototype ps
+                   return ((tx : rndTxs) ++ rest)
+           expandedTxs <- expandPrototype seqPrototype
+           corpusSet <- liftIO $ readIORef env.corpusRef
+           wid <- gets (.workerId)
+           -- Select a prefix from the existing corpus
+           -- Special handling for worker 0: always use empty prefix (position 0)
+           prefix <- if Set.null corpusSet || wid == 0
+                     then pure []
+                     else do
+                       -- Pick a random sequence from corpus
+                       idx <- getRandomR (0, Set.size corpusSet - 1)
+                       let (_, cTxs) = Set.elemAt idx corpusSet
+                       let middleLen = length expandedTxs
+                       let maxPrefix = seqLen - middleLen
+                       if maxPrefix <= 0
+                         then pure []
+                         else do
+                           -- Take a random prefix length
+                           k <- getRandomR (0, min (length cTxs) maxPrefix)
+                           pure (take k cTxs)
+           let combined = prefix ++ expandedTxs
+           let len = length combined
+           -- Pad with random transactions if sequence is too short
+           if len < seqLen
+             then do
+               paddingTxs <- replicateM (seqLen - len) (genTx world deployedContracts)
+               pure (combined ++ paddingTxs)
+             else
+               pure (take seqLen combined)
+    -- | Generate a sequence of transactions using standard fuzzing techniques
+    genStandardSeq
+      :: (MonadRandom m, MonadReader Env m, MonadState WorkerState m, MonadIO m)
+      => Map (Expr 'EAddr) Contract
+      -> m [Tx]
+    genStandardSeq deployedContracts = do
+           env <- ask
+           let world = env.world
+               mutConsts = env.cfg.campaignConf.mutConsts
+               seqLen = env.cfg.campaignConf.seqLen
+           -- 3. Standard fuzzing behavior (no prioritized sequence selected)
+           -- Generate new random transactions
+           randTxs <- replicateM seqLen (genTx world deployedContracts)
+           -- Generate a random mutator
+           cmut <- if seqLen == 1 then seqMutatorsStateless (fromConsts mutConsts)
+                                  else seqMutatorsStateful (fromConsts mutConsts)
+           -- Fetch the mutator
+           let mut = getCorpusMutation cmut
+           corpus <- liftIO $ readIORef env.corpusRef
+           if null corpus
+             then pure randTxs -- Use the generated random transactions
+             else mut seqLen corpus randTxs -- Apply the mutator

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Agentic refactoring and full MCP support #1502

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Agentic refactoring and full MCP support #1502

Are you sure you want to change the base?

Uh oh!

Agentic refactoring and full MCP support #1502

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!