From dfc77babca56d202c0750f0f03e09d254cff0ebe Mon Sep 17 00:00:00 2001 From: Mouradif Date: Tue, 17 Mar 2026 16:16:58 +0100 Subject: [PATCH] feat: add fork block number and selector on fork exit --- internal/adapters/anvil/manager.go | 3 + internal/cli/fork.go | 72 ++++++++++-- internal/domain/anvil.go | 15 +-- internal/domain/fork.go | 12 ++ internal/usecase/enter_fork.go | 24 ++-- test/integration/fork_test.go | 175 +++++++++++++++++++++++++++++ 6 files changed, 275 insertions(+), 26 deletions(-) diff --git a/internal/adapters/anvil/manager.go b/internal/adapters/anvil/manager.go index 026237c..eb9f517 100644 --- a/internal/adapters/anvil/manager.go +++ b/internal/adapters/anvil/manager.go @@ -207,6 +207,9 @@ func buildAnvilArgs(instance *domain.AnvilInstance) []string { if instance.ForkURL != "" { args = append(args, "--fork-url", instance.ForkURL) } + if instance.ForkBlockNumber > 0 { + args = append(args, "--fork-block-number", strconv.FormatUint(instance.ForkBlockNumber, 10)) + } return args } diff --git a/internal/cli/fork.go b/internal/cli/fork.go index 8244cff..75003e3 100644 --- a/internal/cli/fork.go +++ b/internal/cli/fork.go @@ -4,7 +4,10 @@ import ( "encoding/json" "fmt" "os" + "sort" + "github.com/fatih/color" + "github.com/manifoldco/promptui" "github.com/spf13/cobra" "github.com/trebuchet-org/treb-cli/internal/cli/render" cfg "github.com/trebuchet-org/treb-cli/internal/config" @@ -48,6 +51,7 @@ Use --url to connect to an already-running Anvil fork instead of starting one lo } cmd.Flags().String("url", "", "Connect to an external Anvil endpoint instead of starting a local fork") + cmd.Flags().Uint64("fork-block-number", 0, "Fork at a specific block number (default: latest)") return cmd } @@ -118,13 +122,16 @@ func runForkEnter(cmd *cobra.Command, args []string) error { return fmt.Errorf("failed to resolve network '%s': %w", network, err) } + forkBlockNumber, _ := cmd.Flags().GetUint64("fork-block-number") + // Execute the use case params := usecase.EnterForkParams{ - Network: network, - RPCURL: resolvedNetwork.RPCURL, - ChainID: resolvedNetwork.ChainID, - EnvVarName: envVarName, - ExternalURL: externalURL, + Network: network, + RPCURL: resolvedNetwork.RPCURL, + ChainID: resolvedNetwork.ChainID, + EnvVarName: envVarName, + ExternalURL: externalURL, + ForkBlockNumber: forkBlockNumber, } result, err := app.EnterFork.Execute(ctx, params) @@ -170,9 +177,31 @@ func runForkExit(cmd *cobra.Command, args []string) error { if len(args) > 0 { network = args[0] } else if !allFlag { - // Use current configured network - if app.Config.Network != nil { - network = app.Config.Network.Name + // No network specified and not --all: check if multiple forks are active + state, err := app.ForkStateStore.Load(ctx) + if err != nil { + return fmt.Errorf("failed to load fork state: %w", err) + } + + switch len(state.Forks) { + case 0: + return fmt.Errorf("no active forks") + case 1: + // Only one fork — use it directly + for name := range state.Forks { + network = name + } + default: + // Multiple forks — prompt user to select + if app.Config.NonInteractive { + return fmt.Errorf("multiple active forks. Specify a network or use --all") + } + + selected, err := selectForkNetwork(state.ActiveNetworks()) + if err != nil { + return err + } + network = selected } } @@ -190,6 +219,33 @@ func runForkExit(cmd *cobra.Command, args []string) error { return renderer.RenderExit(result) } +// selectForkNetwork prompts the user to select a fork network from a list +func selectForkNetwork(networks []string) (string, error) { + sort.Strings(networks) + + templates := &promptui.SelectTemplates{ + Label: "{{ . }}", + Active: "▸ {{ . | cyan }}", + Inactive: " {{ . | faint }}", + Selected: "✓ {{ . | green }}", + Help: color.New(color.FgYellow).Sprint("Use arrow keys to navigate, Enter to select"), + } + + prompt := promptui.Select{ + Label: "Multiple active forks. Which one do you want to exit?", + Items: networks, + Templates: templates, + Size: 10, + } + + _, selected, err := prompt.Run() + if err != nil { + return "", fmt.Errorf("selection cancelled: %w", err) + } + + return selected, nil +} + // newForkRevertCmd creates the fork revert subcommand func newForkRevertCmd() *cobra.Command { cmd := &cobra.Command{ diff --git a/internal/domain/anvil.go b/internal/domain/anvil.go index 63e2cec..e57b9a3 100644 --- a/internal/domain/anvil.go +++ b/internal/domain/anvil.go @@ -2,13 +2,14 @@ package domain // AnvilInstance represents a local anvil node instance type AnvilInstance struct { - Name string `json:"name"` - Port string `json:"port"` - ChainID string `json:"chainId,omitempty"` - ForkURL string `json:"forkUrl,omitempty"` - RPCURL string `json:"rpcUrl,omitempty"` // Full RPC endpoint URL (used for external forks) - PidFile string `json:"pidFile"` - LogFile string `json:"logFile"` + Name string `json:"name"` + Port string `json:"port"` + ChainID string `json:"chainId,omitempty"` + ForkURL string `json:"forkUrl,omitempty"` + ForkBlockNumber uint64 `json:"forkBlockNumber,omitempty"` // Fork at a specific block number (0 = latest) + RPCURL string `json:"rpcUrl,omitempty"` // Full RPC endpoint URL (used for external forks) + PidFile string `json:"pidFile"` + LogFile string `json:"logFile"` } // AnvilStatus represents the status of an anvil instance diff --git a/internal/domain/fork.go b/internal/domain/fork.go index b45a81c..c1e5994 100644 --- a/internal/domain/fork.go +++ b/internal/domain/fork.go @@ -57,6 +57,18 @@ func (s *ForkState) GetActiveFork(network string) *ForkEntry { return s.Forks[network] } +// ActiveNetworks returns a list of all active fork network names +func (s *ForkState) ActiveNetworks() []string { + if s == nil || s.Forks == nil { + return nil + } + networks := make([]string, 0, len(s.Forks)) + for name := range s.Forks { + networks = append(networks, name) + } + return networks +} + // AnvilInstance returns an AnvilInstance for this fork entry. // For external forks, RPCURL is set to the full fork URL so RPC calls // go to the correct endpoint instead of constructing http://localhost:. diff --git a/internal/usecase/enter_fork.go b/internal/usecase/enter_fork.go index 02f396e..9c47d06 100644 --- a/internal/usecase/enter_fork.go +++ b/internal/usecase/enter_fork.go @@ -46,11 +46,12 @@ func NewEnterFork( // EnterForkParams contains parameters for entering fork mode type EnterForkParams struct { - Network string // network name from foundry.toml - RPCURL string // resolved RPC URL (after env var expansion) - ChainID uint64 // chain ID - EnvVarName string // env var name that foundry.toml uses for the RPC endpoint - ExternalURL string // optional external Anvil endpoint URL (skips local Anvil startup) + Network string // network name from foundry.toml + RPCURL string // resolved RPC URL (after env var expansion) + ChainID uint64 // chain ID + EnvVarName string // env var name that foundry.toml uses for the RPC endpoint + ExternalURL string // optional external Anvil endpoint URL (skips local Anvil startup) + ForkBlockNumber uint64 // optional block number to fork at (0 = latest) } // EnterForkResult contains the result of entering fork mode @@ -144,12 +145,13 @@ func (uc *EnterFork) executeLocal(ctx context.Context, state *domain.ForkState, } instance := &domain.AnvilInstance{ - Name: fmt.Sprintf("fork-%s", params.Network), - Port: fmt.Sprintf("%d", port), - ChainID: fmt.Sprintf("%d", params.ChainID), - ForkURL: params.RPCURL, - PidFile: filepath.Join(privDir, fmt.Sprintf("fork-%s.pid", params.Network)), - LogFile: filepath.Join(privDir, fmt.Sprintf("fork-%s.log", params.Network)), + Name: fmt.Sprintf("fork-%s", params.Network), + Port: fmt.Sprintf("%d", port), + ChainID: fmt.Sprintf("%d", params.ChainID), + ForkURL: params.RPCURL, + ForkBlockNumber: params.ForkBlockNumber, + PidFile: filepath.Join(privDir, fmt.Sprintf("fork-%s.pid", params.Network)), + LogFile: filepath.Join(privDir, fmt.Sprintf("fork-%s.log", params.Network)), } // Start anvil (includes CreateX deployment) diff --git a/test/integration/fork_test.go b/test/integration/fork_test.go index ac3a733..ce2d31a 100644 --- a/test/integration/fork_test.go +++ b/test/integration/fork_test.go @@ -1685,6 +1685,181 @@ func TestForkExternalExit(t *testing.T) { RunIntegrationTests(t, tests) } +func TestForkEnterBlockNumber(t *testing.T) { + tests := []IntegrationTest{ + { + Name: "fork_enter_with_block_number", + PreSetup: func(t *testing.T, ctx *helpers.TestContext) { + setupForkEnvVars(t, ctx) + }, + SetupCmds: [][]string{ + s("config set network anvil-31337"), + }, + TestCmds: [][]string{}, + SkipGolden: true, + PostTest: func(t *testing.T, ctx *helpers.TestContext, _ string) { + // First, mine a few blocks on the test anvil so we have a known block to fork at + node := ctx.AnvilNodes["anvil-31337"] + require.NotNil(t, node) + mineBlocks(t, node.URL, 10) + + // Record block number before fork + baseBlock := ethBlockNumber(t, node.URL) + require.GreaterOrEqual(t, baseBlock, uint64(10)) + + // Fork at a specific block + forkBlock := baseBlock - 5 + output, err := ctx.TrebContext.Treb("fork", "enter", "anvil-31337", + "--fork-block-number", strconv.FormatUint(forkBlock, 10)) + require.NoError(t, err, "fork enter with --fork-block-number should succeed") + defer cleanupForkAnvil(t, ctx, "anvil-31337") + + assert.Contains(t, output, "Fork mode entered") + + // Verify fork state exists and anvil is running + state := readForkState(t, ctx) + fork := state.Forks["anvil-31337"] + require.NotNil(t, fork, "fork entry should exist") + assert.True(t, isProcessAlive(fork.AnvilPID), "fork anvil should be running") + + // Verify the fork anvil's block number matches the fork block + blockNum := ethBlockNumber(t, fork.ForkURL) + assert.Equal(t, forkBlock, blockNum, + "fork anvil should be at the specified fork-block-number") + }, + }, + } + + RunIntegrationTests(t, tests) +} + +// mineBlocks mines n empty blocks on an anvil instance via the evm_mine RPC +func mineBlocks(t *testing.T, rpcURL string, n int) { + t.Helper() + for range n { + reqBody, err := json.Marshal(map[string]interface{}{ + "jsonrpc": "2.0", + "method": "evm_mine", + "params": []interface{}{}, + "id": 1, + }) + require.NoError(t, err) + + resp, err := http.Post(rpcURL, "application/json", bytes.NewBuffer(reqBody)) //nolint:gosec + require.NoError(t, err) + resp.Body.Close() + } +} + +func TestForkExitMultipleForks(t *testing.T) { + tests := []IntegrationTest{ + { + Name: "fork_exit_no_network_single_fork_exits_it", + PreSetup: func(t *testing.T, ctx *helpers.TestContext) { + setupForkEnvVars(t, ctx) + }, + SetupCmds: [][]string{ + s("config set network anvil-31337"), + {"fork", "enter", "anvil-31337"}, + }, + TestCmds: [][]string{ + // No network argument — should auto-select the only active fork + {"fork", "exit"}, + }, + SkipGolden: true, + PostTest: func(t *testing.T, ctx *helpers.TestContext, output string) { + assert.Contains(t, output, "Fork mode exited") + assert.Contains(t, output, "anvil-31337") + + // Verify fork state file is gone + workDir := ctx.TrebContext.GetWorkDir() + statePath := filepath.Join(workDir, ".treb", "priv", "fork-state.json") + _, err := os.Stat(statePath) + assert.True(t, os.IsNotExist(err), "fork-state.json should be deleted") + }, + }, + { + Name: "fork_exit_no_network_multiple_forks_errors_non_interactive", + PreSetup: func(t *testing.T, ctx *helpers.TestContext) { + setupForkEnvVars(t, ctx) + }, + SetupCmds: [][]string{ + s("config set network anvil-31337"), + {"fork", "enter", "anvil-31337"}, + {"fork", "enter", "anvil-31338"}, + }, + TestCmds: [][]string{ + // No network argument, multiple forks, non-interactive — should error + {"fork", "exit"}, + }, + ExpectErr: true, + SkipGolden: true, + PostTest: func(t *testing.T, ctx *helpers.TestContext, output string) { + defer cleanupForkAnvil(t, ctx, "anvil-31337") + defer cleanupForkAnvil(t, ctx, "anvil-31338") + + assert.Contains(t, output, "multiple active forks") + }, + }, + { + Name: "fork_exit_all_exits_multiple_forks", + PreSetup: func(t *testing.T, ctx *helpers.TestContext) { + setupForkEnvVars(t, ctx) + }, + SetupCmds: [][]string{ + s("config set network anvil-31337"), + {"fork", "enter", "anvil-31337"}, + {"fork", "enter", "anvil-31338"}, + }, + TestCmds: [][]string{ + {"fork", "exit", "--all"}, + }, + SkipGolden: true, + PostTest: func(t *testing.T, ctx *helpers.TestContext, output string) { + assert.Contains(t, output, "2 fork(s) exited") + + // Verify fork state file is gone + workDir := ctx.TrebContext.GetWorkDir() + statePath := filepath.Join(workDir, ".treb", "priv", "fork-state.json") + _, err := os.Stat(statePath) + assert.True(t, os.IsNotExist(err), "fork-state.json should be deleted") + }, + }, + } + + RunIntegrationTests(t, tests) +} + +// ethBlockNumber makes an eth_blockNumber RPC call and returns the block number +func ethBlockNumber(t *testing.T, rpcURL string) uint64 { + t.Helper() + + reqBody, err := json.Marshal(map[string]interface{}{ + "jsonrpc": "2.0", + "method": "eth_blockNumber", + "params": []interface{}{}, + "id": 1, + }) + require.NoError(t, err) + + resp, err := http.Post(rpcURL, "application/json", bytes.NewBuffer(reqBody)) //nolint:gosec + require.NoError(t, err) + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + var rpcResp struct { + Result string `json:"result"` + } + require.NoError(t, json.Unmarshal(body, &rpcResp)) + + blockNum, err := strconv.ParseUint(strings.TrimPrefix(rpcResp.Result, "0x"), 16, 64) + require.NoError(t, err) + + return blockNum +} + func TestForkExternalStatus(t *testing.T) { tests := []IntegrationTest{ {