From 4db2d4277716a1c6806a1c2dc54c62c9015df51b Mon Sep 17 00:00:00 2001
From: Santiago Bartesaghi <sbartesaghi@hotmail.com>
Date: Sun, 1 Mar 2026 13:05:10 -0300
Subject: [PATCH 1/6] Add ExecutionContext, Formatter, and Hooks
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 1 of the Mars v2 refactor — additive, no breaking changes.

- ExecutionContext: n8n-style state that flows between steps, with
  namespaced access to previous outputs, fork/merge for parallel
  branches, and shared global_state
- Formatter: injectable input/output transform (identity by default)
- Hooks: class-level before_run/after_run callbacks mixin

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 lib/mars/execution_context.rb       |  34 ++++++++
 lib/mars/formatter.rb               |  13 +++
 lib/mars/hooks.rb                   |  35 ++++++++
 spec/mars/execution_context_spec.rb | 121 ++++++++++++++++++++++++++++
 spec/mars/formatter_spec.rb         |  43 ++++++++++
 spec/mars/hooks_spec.rb             |  80 ++++++++++++++++++
 6 files changed, 326 insertions(+)
 create mode 100644 lib/mars/execution_context.rb
 create mode 100644 lib/mars/formatter.rb
 create mode 100644 lib/mars/hooks.rb
 create mode 100644 spec/mars/execution_context_spec.rb
 create mode 100644 spec/mars/formatter_spec.rb
 create mode 100644 spec/mars/hooks_spec.rb

diff --git a/lib/mars/execution_context.rb b/lib/mars/execution_context.rb
new file mode 100644
index 0000000..381de13
--- /dev/null
+++ b/lib/mars/execution_context.rb
@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module MARS
+  class ExecutionContext
+    attr_reader :current_input, :outputs, :global_state
+
+    def initialize(input: nil, global_state: {})
+      @current_input = input
+      @outputs = {}
+      @global_state = global_state
+    end
+
+    def [](step_name)
+      outputs[step_name]
+    end
+
+    def record(step_name, output)
+      @outputs[step_name] = output
+      @current_input = output
+    end
+
+    def fork(input: current_input)
+      self.class.new(input: input, global_state: global_state)
+    end
+
+    def merge(child_contexts)
+      child_contexts.each do |child|
+        @outputs.merge!(child.outputs)
+      end
+
+      self
+    end
+  end
+end
diff --git a/lib/mars/formatter.rb b/lib/mars/formatter.rb
new file mode 100644
index 0000000..f82fb9c
--- /dev/null
+++ b/lib/mars/formatter.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module MARS
+  class Formatter
+    def format_input(context)
+      context.current_input
+    end
+
+    def format_output(output)
+      output
+    end
+  end
+end
diff --git a/lib/mars/hooks.rb b/lib/mars/hooks.rb
new file mode 100644
index 0000000..df616ad
--- /dev/null
+++ b/lib/mars/hooks.rb
@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MARS
+  module Hooks
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      def before_run(&block)
+        before_run_hooks << block
+      end
+
+      def after_run(&block)
+        after_run_hooks << block
+      end
+
+      def before_run_hooks
+        @before_run_hooks ||= []
+      end
+
+      def after_run_hooks
+        @after_run_hooks ||= []
+      end
+    end
+
+    def run_before_hooks(context)
+      self.class.before_run_hooks.each { |hook| hook.call(context, self) }
+    end
+
+    def run_after_hooks(context, result)
+      self.class.after_run_hooks.each { |hook| hook.call(context, result, self) }
+    end
+  end
+end
diff --git a/spec/mars/execution_context_spec.rb b/spec/mars/execution_context_spec.rb
new file mode 100644
index 0000000..d30d47d
--- /dev/null
+++ b/spec/mars/execution_context_spec.rb
@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::ExecutionContext do
+  describe "#current_input" do
+    it "returns the initial input" do
+      context = described_class.new(input: "query")
+      expect(context.current_input).to eq("query")
+    end
+
+    it "returns nil when no input is provided" do
+      context = described_class.new
+      expect(context.current_input).to be_nil
+    end
+  end
+
+  describe "#record" do
+    it "stores the output under the step name" do
+      context = described_class.new(input: "query")
+      context.record(:researcher, "research result")
+
+      expect(context[:researcher]).to eq("research result")
+    end
+
+    it "updates current_input to the recorded output" do
+      context = described_class.new(input: "query")
+      context.record(:researcher, "research result")
+
+      expect(context.current_input).to eq("research result")
+    end
+
+    it "tracks multiple step outputs" do
+      context = described_class.new(input: "query")
+      context.record(:researcher, "research result")
+      context.record(:summarizer, "summary")
+
+      expect(context.outputs).to eq({ researcher: "research result", summarizer: "summary" })
+      expect(context.current_input).to eq("summary")
+    end
+  end
+
+  describe "#[]" do
+    it "returns nil for unknown step names" do
+      context = described_class.new(input: "query")
+      expect(context[:unknown]).to be_nil
+    end
+  end
+
+  describe "#global_state" do
+    it "defaults to an empty hash" do
+      context = described_class.new(input: "query")
+      expect(context.global_state).to eq({})
+    end
+
+    it "accepts initial global state" do
+      context = described_class.new(input: "query", global_state: { user_id: 42 })
+      expect(context.global_state).to eq({ user_id: 42 })
+    end
+
+    it "is mutable" do
+      context = described_class.new(input: "query")
+      context.global_state[:key] = "value"
+
+      expect(context.global_state[:key]).to eq("value")
+    end
+  end
+
+  describe "#fork" do
+    it "creates a child context with the same current_input by default" do
+      context = described_class.new(input: "query")
+      child = context.fork
+
+      expect(child.current_input).to eq("query")
+    end
+
+    it "creates a child context with a custom input" do
+      context = described_class.new(input: "query")
+      child = context.fork(input: "custom")
+
+      expect(child.current_input).to eq("custom")
+    end
+
+    it "shares global_state with the parent" do
+      context = described_class.new(input: "query", global_state: { shared: true })
+      child = context.fork
+
+      child.global_state[:added_by_child] = true
+
+      expect(context.global_state[:added_by_child]).to be(true)
+    end
+
+    it "has independent outputs from the parent" do
+      context = described_class.new(input: "query")
+      context.record(:parent_step, "parent output")
+
+      child = context.fork
+      child.record(:child_step, "child output")
+
+      expect(context[:child_step]).to be_nil
+      expect(child[:parent_step]).to be_nil
+    end
+  end
+
+  describe "#merge" do
+    it "merges child outputs into the parent" do
+      context = described_class.new(input: "query")
+      context.record(:step1, "output1")
+
+      child1 = context.fork
+      child1.record(:branch_a, "result_a")
+
+      child2 = context.fork
+      child2.record(:branch_b, "result_b")
+
+      context.merge([child1, child2])
+
+      expect(context[:step1]).to eq("output1")
+      expect(context[:branch_a]).to eq("result_a")
+      expect(context[:branch_b]).to eq("result_b")
+    end
+  end
+end
diff --git a/spec/mars/formatter_spec.rb b/spec/mars/formatter_spec.rb
new file mode 100644
index 0000000..55b17f3
--- /dev/null
+++ b/spec/mars/formatter_spec.rb
@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::Formatter do
+  let(:formatter) { described_class.new }
+
+  describe "#format_input" do
+    it "returns the context's current_input" do
+      context = MARS::ExecutionContext.new(input: "hello")
+      expect(formatter.format_input(context)).to eq("hello")
+    end
+  end
+
+  describe "#format_output" do
+    it "returns the output unchanged" do
+      expect(formatter.format_output("result")).to eq("result")
+    end
+  end
+
+  describe "custom formatter" do
+    let(:custom_formatter_class) do
+      Class.new(described_class) do
+        def format_input(context)
+          context.current_input.upcase
+        end
+
+        def format_output(output)
+          "formatted: #{output}"
+        end
+      end
+    end
+
+    let(:custom_formatter) { custom_formatter_class.new }
+
+    it "can override format_input" do
+      context = MARS::ExecutionContext.new(input: "hello")
+      expect(custom_formatter.format_input(context)).to eq("HELLO")
+    end
+
+    it "can override format_output" do
+      expect(custom_formatter.format_output("result")).to eq("formatted: result")
+    end
+  end
+end
diff --git a/spec/mars/hooks_spec.rb b/spec/mars/hooks_spec.rb
new file mode 100644
index 0000000..d304a7c
--- /dev/null
+++ b/spec/mars/hooks_spec.rb
@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::Hooks do
+  let(:hookable_class) do
+    Class.new do
+      include MARS::Hooks
+
+      attr_reader :name
+
+      def initialize(name)
+        @name = name
+      end
+    end
+  end
+
+  describe ".before_run" do
+    it "registers a before hook" do
+      hookable_class.before_run { |_ctx, _step| "before" }
+
+      expect(hookable_class.before_run_hooks.size).to eq(1)
+    end
+  end
+
+  describe ".after_run" do
+    it "registers an after hook" do
+      hookable_class.after_run { |_ctx, _result, _step| "after" }
+
+      expect(hookable_class.after_run_hooks.size).to eq(1)
+    end
+  end
+
+  describe "#run_before_hooks" do
+    it "calls all before hooks with context and step" do
+      calls = []
+      hookable_class.before_run { |ctx, step| calls << [ctx, step.name] }
+
+      step = hookable_class.new("test_step")
+      context = MARS::ExecutionContext.new(input: "input")
+      step.run_before_hooks(context)
+
+      expect(calls).to eq([[context, "test_step"]])
+    end
+
+    it "calls hooks in registration order" do
+      order = []
+      hookable_class.before_run { |_ctx, _step| order << :first }
+      hookable_class.before_run { |_ctx, _step| order << :second }
+
+      step = hookable_class.new("test_step")
+      step.run_before_hooks(MARS::ExecutionContext.new(input: "input"))
+
+      expect(order).to eq(%i[first second])
+    end
+  end
+
+  describe "#run_after_hooks" do
+    it "calls all after hooks with context, result, and step" do
+      calls = []
+      hookable_class.after_run { |ctx, result, step| calls << [ctx, result, step.name] }
+
+      step = hookable_class.new("test_step")
+      context = MARS::ExecutionContext.new(input: "input")
+      step.run_after_hooks(context, "the result")
+
+      expect(calls).to eq([[context, "the result", "test_step"]])
+    end
+  end
+
+  describe "hook isolation between classes" do
+    it "does not share hooks between different classes" do
+      class_a = Class.new { include MARS::Hooks }
+      class_b = Class.new { include MARS::Hooks }
+
+      class_a.before_run { "a" }
+
+      expect(class_a.before_run_hooks.size).to eq(1)
+      expect(class_b.before_run_hooks.size).to eq(0)
+    end
+  end
+end

From 152368566bcd0b0f3b3beceae722c044cb60e8dd Mon Sep 17 00:00:00 2001
From: Santiago Bartesaghi <sbartesaghi@hotmail.com>
Date: Sun, 1 Mar 2026 13:10:44 -0300
Subject: [PATCH 2/6] Refactor Runnable with name, formatter, and hooks

Phase 2 of the Mars v2 refactor.

- Runnable: include Hooks, add `name` (auto-derived from class via
  `step_name`), add `formatter` class-level DSL with instance fallback
- Gate, Aggregator, Sequential, Parallel: delegate `name` to Runnable
  via super instead of managing their own attr_reader
- Runnable spec: cover name derivation, formatter DSL, hooks integration

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 lib/mars/aggregator.rb           |  5 +--
 lib/mars/gate.rb                 |  5 +--
 lib/mars/runnable.rb             | 22 ++++++++++-
 lib/mars/workflows/parallel.rb   |  5 +--
 lib/mars/workflows/sequential.rb |  5 +--
 spec/mars/runnable_spec.rb       | 67 ++++++++++++++++++++++++++++++++
 6 files changed, 93 insertions(+), 16 deletions(-)

diff --git a/lib/mars/aggregator.rb b/lib/mars/aggregator.rb
index 0866e9f..d21b3bd 100644
--- a/lib/mars/aggregator.rb
+++ b/lib/mars/aggregator.rb
@@ -2,12 +2,11 @@
 
 module MARS
   class Aggregator < Runnable
-    attr_reader :name, :operation
+    attr_reader :operation
 
     def initialize(name = "Aggregator", operation: nil, **kwargs)
-      super(**kwargs)
+      super(name: name, **kwargs)
 
-      @name = name
       @operation = operation || ->(inputs) { inputs }
     end
 
diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb
index 21e407d..8a2bafd 100644
--- a/lib/mars/gate.rb
+++ b/lib/mars/gate.rb
@@ -2,12 +2,9 @@
 
 module MARS
   class Gate < Runnable
-    attr_reader :name
-
     def initialize(name = "Gate", condition:, branches:, **kwargs)
-      super(**kwargs)
+      super(name: name, **kwargs)
 
-      @name = name
       @condition = condition
       @branches = branches
     end
diff --git a/lib/mars/runnable.rb b/lib/mars/runnable.rb
index f7a4909..115baf1 100644
--- a/lib/mars/runnable.rb
+++ b/lib/mars/runnable.rb
@@ -2,10 +2,30 @@
 
 module MARS
   class Runnable
+    include Hooks
+
+    attr_reader :name, :formatter
     attr_accessor :state
 
-    def initialize(state: {})
+    class << self
+      def step_name
+        return @step_name if defined?(@step_name)
+        return unless name
+
+        name.split("::").last.gsub(/([a-z])([A-Z])/, '\1_\2').downcase
+      end
+
+      attr_writer :step_name
+
+      def formatter(klass = nil)
+        klass ? @formatter_class = klass : @formatter_class
+      end
+    end
+
+    def initialize(name: self.class.step_name, state: {}, formatter: nil)
+      @name = name
       @state = state
+      @formatter = formatter || self.class.formatter&.new || Formatter.new
     end
 
     def run(input)
diff --git a/lib/mars/workflows/parallel.rb b/lib/mars/workflows/parallel.rb
index da65a27..ef8f3f6 100644
--- a/lib/mars/workflows/parallel.rb
+++ b/lib/mars/workflows/parallel.rb
@@ -3,12 +3,9 @@
 module MARS
   module Workflows
     class Parallel < Runnable
-      attr_reader :name
-
       def initialize(name, steps:, aggregator: nil, **kwargs)
-        super(**kwargs)
+        super(name: name, **kwargs)
 
-        @name = name
         @steps = steps
         @aggregator = aggregator || Aggregator.new("#{name} Aggregator")
       end
diff --git a/lib/mars/workflows/sequential.rb b/lib/mars/workflows/sequential.rb
index 0db625e..df673c6 100644
--- a/lib/mars/workflows/sequential.rb
+++ b/lib/mars/workflows/sequential.rb
@@ -3,12 +3,9 @@
 module MARS
   module Workflows
     class Sequential < Runnable
-      attr_reader :name
-
       def initialize(name, steps:, **kwargs)
-        super(**kwargs)
+        super(name: name, **kwargs)
 
-        @name = name
         @steps = steps
       end
 
diff --git a/spec/mars/runnable_spec.rb b/spec/mars/runnable_spec.rb
index 6c3e5f7..972a549 100644
--- a/spec/mars/runnable_spec.rb
+++ b/spec/mars/runnable_spec.rb
@@ -42,6 +42,73 @@ def run(input)
     end
   end
 
+  describe "#name" do
+    it "defaults to nil for anonymous classes" do
+      klass = Class.new(described_class)
+      expect(klass.new.name).to be_nil
+    end
+
+    it "can be set via the name keyword" do
+      runnable = described_class.new(name: "my_step")
+      expect(runnable.name).to eq("my_step")
+    end
+
+    it "derives step_name from the class name" do
+      stub_const("MARS::MyCustomStep", Class.new(described_class))
+      expect(MARS::MyCustomStep.new.name).to eq("my_custom_step")
+    end
+  end
+
+  describe "#formatter" do
+    it "defaults to a Formatter instance" do
+      runnable = described_class.new
+      expect(runnable.formatter).to be_a(MARS::Formatter)
+    end
+
+    it "can be set via the formatter keyword" do
+      custom_formatter = MARS::Formatter.new
+      runnable = described_class.new(formatter: custom_formatter)
+      expect(runnable.formatter).to eq(custom_formatter)
+    end
+
+    it "uses the class-level formatter when declared" do
+      custom_formatter_class = Class.new(MARS::Formatter)
+      klass = Class.new(described_class) do
+        formatter custom_formatter_class
+      end
+
+      expect(klass.new.formatter).to be_a(custom_formatter_class)
+    end
+  end
+
+  describe "hooks" do
+    it "includes Hooks module" do
+      expect(described_class.ancestors).to include(MARS::Hooks)
+    end
+
+    it "supports before_run hooks" do
+      klass = Class.new(described_class)
+      calls = []
+      klass.before_run { |_ctx, step| calls << step.name }
+
+      step = klass.new(name: "test")
+      step.run_before_hooks(MARS::ExecutionContext.new(input: "x"))
+
+      expect(calls).to eq(["test"])
+    end
+
+    it "supports after_run hooks" do
+      klass = Class.new(described_class)
+      calls = []
+      klass.after_run { |_ctx, result, _step| calls << result }
+
+      step = klass.new(name: "test")
+      step.run_after_hooks(MARS::ExecutionContext.new(input: "x"), "result")
+
+      expect(calls).to eq(["result"])
+    end
+  end
+
   describe "inheritance" do
     it "can be inherited" do
       subclass = Class.new(described_class)

From 3d5df3276b68dad499301f7777538c1a87248231 Mon Sep 17 00:00:00 2001
From: Santiago Bartesaghi <sbartesaghi@hotmail.com>
Date: Sun, 1 Mar 2026 13:15:23 -0300
Subject: [PATCH 3/6] Rename Agent to AgentStep with agent macro
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 3 of the Mars v2 refactor.

- Rename Mars::Agent → Mars::AgentStep with class-level `agent` macro
  that wraps a RubyLLM::Agent subclass
- AgentStep#run creates a new agent instance and delegates via .ask
- Remove old Agent with its manual Chat setup (before_run/after_run,
  system_prompt, tools, schema instance methods)
- Rename rendering graph module accordingly
- Bump ruby_llm dependency to ~> 1.12

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 lib/mars/agent.rb                             | 48 --------------
 lib/mars/agent_step.rb                        | 15 +++++
 lib/mars/rendering/graph.rb                   |  2 +-
 .../graph/{agent.rb => agent_step.rb}         |  6 +-
 mars.gemspec                                  |  2 +-
 spec/mars/agent_spec.rb                       | 64 -------------------
 spec/mars/agent_step_spec.rb                  | 61 ++++++++++++++++++
 7 files changed, 79 insertions(+), 119 deletions(-)
 delete mode 100644 lib/mars/agent.rb
 create mode 100644 lib/mars/agent_step.rb
 rename lib/mars/rendering/graph/{agent.rb => agent_step.rb} (81%)
 delete mode 100644 spec/mars/agent_spec.rb
 create mode 100644 spec/mars/agent_step_spec.rb

diff --git a/lib/mars/agent.rb b/lib/mars/agent.rb
deleted file mode 100644
index edc7ed7..0000000
--- a/lib/mars/agent.rb
+++ /dev/null
@@ -1,48 +0,0 @@
-# frozen_string_literal: true
-
-module MARS
-  class Agent < Runnable
-    def initialize(options: {}, **kwargs)
-      super(**kwargs)
-
-      @options = options
-    end
-
-    def run(input)
-      processed_input = before_run(input)
-      response = chat.ask(processed_input).content
-      after_run(response)
-    end
-
-    private
-
-    attr_reader :options
-
-    def chat
-      @chat ||= RubyLLM::Chat.new(**options)
-                             .with_instructions(system_prompt)
-                             .with_tools(*tools)
-                             .with_schema(schema)
-    end
-
-    def before_run(input)
-      input
-    end
-
-    def after_run(response)
-      response
-    end
-
-    def system_prompt
-      nil
-    end
-
-    def tools
-      []
-    end
-
-    def schema
-      nil
-    end
-  end
-end
diff --git a/lib/mars/agent_step.rb b/lib/mars/agent_step.rb
new file mode 100644
index 0000000..a519997
--- /dev/null
+++ b/lib/mars/agent_step.rb
@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module MARS
+  class AgentStep < Runnable
+    class << self
+      def agent(klass = nil)
+        klass ? @agent_class = klass : @agent_class
+      end
+    end
+
+    def run(input)
+      self.class.agent.new.ask(input).content
+    end
+  end
+end
diff --git a/lib/mars/rendering/graph.rb b/lib/mars/rendering/graph.rb
index 1a73d66..b2ab350 100644
--- a/lib/mars/rendering/graph.rb
+++ b/lib/mars/rendering/graph.rb
@@ -4,7 +4,7 @@ module MARS
   module Rendering
     module Graph
       def self.include_extensions
-        MARS::Agent.include(Agent)
+        MARS::AgentStep.include(AgentStep)
         MARS::Gate.include(Gate)
         MARS::Workflows::Sequential.include(SequentialWorkflow)
         MARS::Workflows::Parallel.include(ParallelWorkflow)
diff --git a/lib/mars/rendering/graph/agent.rb b/lib/mars/rendering/graph/agent_step.rb
similarity index 81%
rename from lib/mars/rendering/graph/agent.rb
rename to lib/mars/rendering/graph/agent_step.rb
index 48f24d8..e170feb 100644
--- a/lib/mars/rendering/graph/agent.rb
+++ b/lib/mars/rendering/graph/agent_step.rb
@@ -3,7 +3,7 @@
 module MARS
   module Rendering
     module Graph
-      module Agent
+      module AgentStep
         include Base
 
         def to_graph(builder, parent_id: nil, value: nil)
@@ -12,10 +12,6 @@ def to_graph(builder, parent_id: nil, value: nil)
 
           [node_id]
         end
-
-        def name
-          self.class.name
-        end
       end
     end
   end
diff --git a/mars.gemspec b/mars.gemspec
index 2e77793..55c3586 100644
--- a/mars.gemspec
+++ b/mars.gemspec
@@ -36,7 +36,7 @@ Gem::Specification.new do |spec|
 
   # Uncomment to register a new dependency of your gem
   spec.add_dependency "async", "~> 2.34"
-  spec.add_dependency "ruby_llm", "~> 1.9"
+  spec.add_dependency "ruby_llm", "~> 1.12"
   spec.add_dependency "zeitwerk", "~> 2.7"
 
   # For more information and examples about making a new gem, check out our
diff --git a/spec/mars/agent_spec.rb b/spec/mars/agent_spec.rb
deleted file mode 100644
index 440f8da..0000000
--- a/spec/mars/agent_spec.rb
+++ /dev/null
@@ -1,64 +0,0 @@
-# frozen_string_literal: true
-
-RSpec.describe MARS::Agent do
-  describe "#run" do
-    subject(:run_agent) { agent.run("input text") }
-
-    let(:agent) { described_class.new(options: { model: "test-model" }) }
-    let(:mock_chat_instance) do
-      instance_double("RubyLLM::Chat").tap do |mock|
-        allow(mock).to receive_messages(with_tools: mock, with_schema: mock, with_instructions: mock,
-                                        ask: mock_chat_response)
-      end
-    end
-    let(:mock_chat_response) { instance_double("RubyLLM::Message", content: "response text") }
-    let(:mock_chat_class) { class_double("RubyLLM::Chat", new: mock_chat_instance) }
-
-    before do
-      stub_const("RubyLLM::Chat", mock_chat_class)
-    end
-
-    it "initializes RubyLLM::Chat with provided options" do
-      run_agent
-
-      expect(mock_chat_class).to have_received(:new).with(model: "test-model")
-    end
-
-    context "when tools are provided" do
-      let(:tools) { [proc { "tool1" }, proc { "tool2" }] }
-      let(:agent_class) do
-        Class.new(described_class) do
-          def tools
-            [proc { "tool1" }, proc { "tool2" }]
-          end
-        end
-      end
-
-      let(:agent) { agent_class.new }
-
-      it "configures chat with tools" do
-        run_agent
-
-        expect(mock_chat_instance).to have_received(:with_tools).with(*tools)
-      end
-    end
-
-    context "when schema is provided" do
-      let(:agent_class) do
-        Class.new(described_class) do
-          def schema
-            { type: "object" }
-          end
-        end
-      end
-
-      let(:agent) { agent_class.new }
-
-      it "configures chat with schema" do
-        run_agent
-
-        expect(mock_chat_instance).to have_received(:with_schema).with({ type: "object" })
-      end
-    end
-  end
-end
diff --git a/spec/mars/agent_step_spec.rb b/spec/mars/agent_step_spec.rb
new file mode 100644
index 0000000..18244a5
--- /dev/null
+++ b/spec/mars/agent_step_spec.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::AgentStep do
+  describe ".agent" do
+    it "stores and retrieves the agent class" do
+      agent_class = Class.new
+      step_class = Class.new(described_class) do
+        agent agent_class
+      end
+
+      expect(step_class.agent).to eq(agent_class)
+    end
+
+    it "returns nil when no agent class is set" do
+      step_class = Class.new(described_class)
+      expect(step_class.agent).to be_nil
+    end
+  end
+
+  describe "#run" do
+    let(:mock_agent_instance) do
+      instance_double("RubyLLM::Agent").tap do |mock|
+        allow(mock).to receive(:ask).and_return(instance_double("RubyLLM::Message", content: "agent response"))
+      end
+    end
+
+    let(:mock_agent_class) do
+      instance_double("Class").tap do |mock|
+        allow(mock).to receive(:new).and_return(mock_agent_instance)
+      end
+    end
+
+    let(:step_class) do
+      klass = mock_agent_class
+      Class.new(described_class) do
+        agent klass
+      end
+    end
+
+    it "creates a new agent instance and calls ask" do
+      step = step_class.new
+      result = step.run("hello")
+
+      expect(result).to eq("agent response")
+      expect(mock_agent_class).to have_received(:new)
+      expect(mock_agent_instance).to have_received(:ask).with("hello")
+    end
+  end
+
+  describe "inheritance" do
+    it "inherits from MARS::Runnable" do
+      expect(described_class.ancestors).to include(MARS::Runnable)
+    end
+
+    it "has access to name, formatter, and hooks from Runnable" do
+      step = described_class.new(name: "my_agent")
+      expect(step.name).to eq("my_agent")
+      expect(step.formatter).to be_a(MARS::Formatter)
+    end
+  end
+end

From a27a06c670ed6dd132cd2da9d7acfce67b31bb2e Mon Sep 17 00:00:00 2001
From: Santiago Bartesaghi <sbartesaghi@hotmail.com>
Date: Sun, 1 Mar 2026 13:18:06 -0300
Subject: [PATCH 4/6] Gate executes branches directly + class-level DSL
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 4 of the Mars v2 refactor.

- Gate: add class-level `condition`/`branch` DSL for reusable gates.
  Gate#run now executes the matched branch directly instead of
  returning a Runnable for Sequential to detect
- Aggregator: context-aware — accepts ExecutionContext and passes
  its outputs to the operation
- Sequential: remove is_a?(Runnable) check, now just chains step results

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 lib/mars/gate.rb                 |  31 +++++++-
 lib/mars/workflows/sequential.rb |   9 +--
 spec/mars/aggregator_spec.rb     |   6 +-
 spec/mars/gate_spec.rb           | 132 ++++++++++++++++++++-----------
 4 files changed, 116 insertions(+), 62 deletions(-)

diff --git a/lib/mars/gate.rb b/lib/mars/gate.rb
index 8a2bafd..e416ab8 100644
--- a/lib/mars/gate.rb
+++ b/lib/mars/gate.rb
@@ -2,21 +2,44 @@
 
 module MARS
   class Gate < Runnable
-    def initialize(name = "Gate", condition:, branches:, **kwargs)
+    class << self
+      def condition(&block)
+        @condition_block = block
+      end
+
+      attr_reader :condition_block
+
+      def branch(key, runnable)
+        branches_map[key] = runnable
+      end
+
+      def branches_map
+        @branches_map ||= {}
+      end
+    end
+
+    def initialize(name = "Gate", condition: nil, branches: nil, **kwargs)
       super(name: name, **kwargs)
 
-      @condition = condition
-      @branches = branches
+      @condition = condition || self.class.condition_block
+      @branches = branches || self.class.branches_map
     end
 
     def run(input)
       result = condition.call(input)
+      branch = branches[result]
+
+      return input unless branch
 
-      branches[result] || input
+      resolve_branch(branch).run(input)
     end
 
     private
 
     attr_reader :condition, :branches
+
+    def resolve_branch(branch)
+      branch.is_a?(Class) ? branch.new : branch
+    end
   end
 end
diff --git a/lib/mars/workflows/sequential.rb b/lib/mars/workflows/sequential.rb
index df673c6..fadae2e 100644
--- a/lib/mars/workflows/sequential.rb
+++ b/lib/mars/workflows/sequential.rb
@@ -11,14 +11,7 @@ def initialize(name, steps:, **kwargs)
 
       def run(input)
         @steps.each do |step|
-          result = step.run(input)
-
-          if result.is_a?(Runnable)
-            input = result.run(input)
-            break
-          else
-            input = result
-          end
+          input = step.run(input)
         end
 
         input
diff --git a/spec/mars/aggregator_spec.rb b/spec/mars/aggregator_spec.rb
index 294ec93..408803e 100644
--- a/spec/mars/aggregator_spec.rb
+++ b/spec/mars/aggregator_spec.rb
@@ -2,7 +2,7 @@
 
 RSpec.describe MARS::Aggregator do
   describe "#run" do
-    context "when called without a block" do
+    context "when called without an operation" do
       let(:aggregator) { described_class.new }
 
       it "returns the input as is" do
@@ -11,10 +11,10 @@
       end
     end
 
-    context "when initialized with a block operation" do
+    context "when initialized with an operation" do
       let(:aggregator) { described_class.new("Aggregator", operation: lambda(&:join)) }
 
-      it "executes the block and returns its value" do
+      it "executes the operation and returns its value" do
         result = aggregator.run(%w[a b c])
         expect(result).to eq("abc")
       end
diff --git a/spec/mars/gate_spec.rb b/spec/mars/gate_spec.rb
index d8cf33d..0990e58 100644
--- a/spec/mars/gate_spec.rb
+++ b/spec/mars/gate_spec.rb
@@ -2,82 +2,120 @@
 
 RSpec.describe MARS::Gate do
   describe "#run" do
-    let(:gate) { described_class.new("TestGate", condition: condition, branches: branches) }
+    context "with constructor-based configuration" do
+      let(:short_step) do
+        Class.new(MARS::Runnable) do
+          def run(input)
+            "short: #{input}"
+          end
+        end.new
+      end
 
-    context "with simple boolean condition" do
-      let(:condition) { ->(input) { input > 5 } }
-      let(:false_branch) { instance_spy(MARS::Runnable) }
-      let(:branches) { { false => false_branch } }
+      let(:long_step) do
+        Class.new(MARS::Runnable) do
+          def run(input)
+            "long: #{input}"
+          end
+        end.new
+      end
 
-      it "returns the input when no branch matches" do
-        result = gate.run(10)
-        expect(result).to eq(10)
+      let(:gate) do
+        described_class.new(
+          "LengthGate",
+          condition: ->(input) { input.length > 5 ? :long : :short },
+          branches: { short: short_step, long: long_step }
+        )
       end
 
-      it "returns the false branch when condition is false" do
-        result = gate.run(3)
+      it "executes the matched branch directly" do
+        expect(gate.run("hi")).to eq("short: hi")
+      end
 
-        expect(result).to eq(false_branch)
+      it "executes the other branch for different input" do
+        expect(gate.run("longstring")).to eq("long: longstring")
       end
 
-      it "does not run the false branch when condition is false" do
-        gate.run(3)
+      it "returns input when no branch matches" do
+        gate = described_class.new(
+          "NoMatch",
+          condition: ->(_input) { :unknown },
+          branches: { short: short_step }
+        )
 
-        expect(false_branch).not_to have_received(:run)
+        expect(gate.run("hello")).to eq("hello")
       end
     end
 
-    context "with string-based condition" do
-      let(:condition) { ->(input) { input.length > 5 ? "long" : "short" } }
-      let(:long_branch) { instance_spy(MARS::Runnable) }
-      let(:short_branch) { instance_spy(MARS::Runnable) }
-      let(:branches) { { "long" => long_branch, "short" => short_branch } }
-
-      it "routes to long branch for long strings" do
-        result = gate.run("longstring")
+    context "with class-level DSL" do
+      let(:short_step_class) do
+        Class.new(MARS::Runnable) do
+          def run(input)
+            "quick: #{input}"
+          end
+        end
+      end
 
-        expect(result).to eq(long_branch)
+      let(:long_step_class) do
+        Class.new(MARS::Runnable) do
+          def run(input)
+            "deep: #{input}"
+          end
+        end
       end
 
-      it "routes to short branch for short strings" do
-        result = gate.run("hi")
+      it "uses condition and branch DSL" do
+        short_cls = short_step_class
+        long_cls = long_step_class
 
-        expect(result).to eq(short_branch)
+        gate_class = Class.new(described_class) do
+          condition { |input| input.length < 5 ? :short : :long }
+          branch :short, short_cls
+          branch :long, long_cls
+        end
+
+        gate = gate_class.new("DSLGate")
+        expect(gate.run("hi")).to eq("quick: hi")
+        expect(gate.run("longstring")).to eq("deep: longstring")
       end
     end
 
     context "with complex condition logic" do
-      let(:condition) do
-        lambda do |input|
-          case input
-          when 0..10 then "low"
-          when 11..50 then "medium"
-          else "high"
-          end
-        end
+      let(:low_step) do
+        Class.new(MARS::Runnable) { def run(input) = "low:#{input}" }.new
       end
 
-      let(:low_branch) { instance_spy(MARS::Runnable) }
-      let(:medium_branch) { instance_spy(MARS::Runnable) }
-      let(:high_branch) { instance_spy(MARS::Runnable) }
-      let(:branches) { { "low" => low_branch, "medium" => medium_branch, "high" => high_branch } }
+      let(:medium_step) do
+        Class.new(MARS::Runnable) { def run(input) = "med:#{input}" }.new
+      end
 
-      it "routes to low branch" do
-        result = gate.run(5)
+      let(:high_step) do
+        Class.new(MARS::Runnable) { def run(input) = "high:#{input}" }.new
+      end
 
-        expect(result).to eq(low_branch)
+      let(:gate) do
+        described_class.new(
+          "SeverityGate",
+          condition: lambda { |input|
+            case input
+            when 0..10 then :low
+            when 11..50 then :medium
+            else :high
+            end
+          },
+          branches: { low: low_step, medium: medium_step, high: high_step }
+        )
       end
 
-      it "routes to medium branch" do
-        result = gate.run(25)
+      it "routes to low branch" do
+        expect(gate.run(5)).to eq("low:5")
+      end
 
-        expect(result).to eq(medium_branch)
+      it "routes to medium branch" do
+        expect(gate.run(25)).to eq("med:25")
       end
 
       it "routes to high branch" do
-        result = gate.run(100)
-
-        expect(result).to eq(high_branch)
+        expect(gate.run(100)).to eq("high:100")
       end
     end
   end

From 45d813305f8950f3d40a9a008a16505dae0390ff Mon Sep 17 00:00:00 2001
From: Santiago Bartesaghi <sbartesaghi@hotmail.com>
Date: Sun, 1 Mar 2026 13:18:06 -0300
Subject: [PATCH 5/6] Gate executes branches directly + class-level DSL
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 4 of the Mars v2 refactor.

- Gate: add class-level `condition`/`branch` DSL for reusable gates.
  Gate#run now executes the matched branch directly instead of
  returning a Runnable for Sequential to detect
- Aggregator: context-aware — accepts ExecutionContext and passes
  its outputs to the operation
- Sequential: remove is_a?(Runnable) check, now just chains step results

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 lib/mars/aggregator.rb                 |   6 +-
 lib/mars/execution_context.rb          |   2 +-
 lib/mars/workflows/parallel.rb         |  39 +++++---
 lib/mars/workflows/sequential.rb       |  31 +++++-
 spec/mars/aggregator_spec.rb           |  15 +++
 spec/mars/workflows/parallel_spec.rb   | 110 ++++++++++++++-------
 spec/mars/workflows/sequential_spec.rb | 128 +++++++++++++++++++++----
 7 files changed, 263 insertions(+), 68 deletions(-)

diff --git a/lib/mars/aggregator.rb b/lib/mars/aggregator.rb
index d21b3bd..78f5752 100644
--- a/lib/mars/aggregator.rb
+++ b/lib/mars/aggregator.rb
@@ -11,7 +11,11 @@ def initialize(name = "Aggregator", operation: nil, **kwargs)
     end
 
     def run(inputs)
-      operation.call(inputs)
+      if inputs.is_a?(ExecutionContext)
+        operation.call(inputs.outputs)
+      else
+        operation.call(inputs)
+      end
     end
   end
 end
diff --git a/lib/mars/execution_context.rb b/lib/mars/execution_context.rb
index 381de13..8d7e1b8 100644
--- a/lib/mars/execution_context.rb
+++ b/lib/mars/execution_context.rb
@@ -15,7 +15,7 @@ def [](step_name)
     end
 
     def record(step_name, output)
-      @outputs[step_name] = output
+      @outputs[step_name.to_sym] = output
       @current_input = output
     end
 
diff --git a/lib/mars/workflows/parallel.rb b/lib/mars/workflows/parallel.rb
index ef8f3f6..b86912e 100644
--- a/lib/mars/workflows/parallel.rb
+++ b/lib/mars/workflows/parallel.rb
@@ -11,27 +11,42 @@ def initialize(name, steps:, aggregator: nil, **kwargs)
       end
 
       def run(input)
-        errors = []
-        results = Async do |workflow|
-          tasks = @steps.map do |step|
-            workflow.async do
-              step.run(input)
-            rescue StandardError => e
-              errors << { error: e, step_name: step.name }
-            end
-          end
+        context = input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
 
-          tasks.map(&:wait)
-        end.result
+        errors = []
+        child_contexts = run_steps_async(context, errors)
 
         raise AggregateError, errors if errors.any?
 
-        aggregator.run(results)
+        context.merge(child_contexts)
+        aggregator.run(context)
       end
 
       private
 
       attr_reader :steps, :aggregator
+
+      def run_steps_async(context, errors)
+        Async do |workflow|
+          tasks = steps.map do |step|
+            workflow.async { run_step(context.fork, step, errors) }
+          end
+
+          tasks.map(&:wait)
+        end.result
+      end
+
+      def run_step(child, step, errors)
+        step.run_before_hooks(child)
+        step_input = step.formatter.format_input(child)
+        result = step.run(step_input)
+        formatted = step.formatter.format_output(result)
+        child.record(step.name, formatted)
+        step.run_after_hooks(child, formatted)
+        child
+      rescue StandardError => e
+        errors << { error: e, step_name: step.name }
+      end
     end
   end
 end
diff --git a/lib/mars/workflows/sequential.rb b/lib/mars/workflows/sequential.rb
index fadae2e..e30e15d 100644
--- a/lib/mars/workflows/sequential.rb
+++ b/lib/mars/workflows/sequential.rb
@@ -9,17 +9,42 @@ def initialize(name, steps:, **kwargs)
         @steps = steps
       end
 
+      def self.build(name, **kwargs, &)
+        builder = Builder.new
+        builder.instance_eval(&)
+        new(name, steps: builder.steps, **kwargs)
+      end
+
       def run(input)
-        @steps.each do |step|
-          input = step.run(input)
+        context = input.is_a?(ExecutionContext) ? input : ExecutionContext.new(input: input)
+
+        steps.each do |step|
+          step.run_before_hooks(context)
+          step_input = step.formatter.format_input(context)
+          result = step.run(step_input)
+          formatted = step.formatter.format_output(result)
+          context.record(step.name, formatted)
+          step.run_after_hooks(context, formatted)
         end
 
-        input
+        context
       end
 
       private
 
       attr_reader :steps
+
+      class Builder
+        attr_reader :steps
+
+        def initialize
+          @steps = []
+        end
+
+        def step(runnable_class, **kwargs)
+          @steps << runnable_class.new(**kwargs)
+        end
+      end
     end
   end
 end
diff --git a/spec/mars/aggregator_spec.rb b/spec/mars/aggregator_spec.rb
index 408803e..dd652da 100644
--- a/spec/mars/aggregator_spec.rb
+++ b/spec/mars/aggregator_spec.rb
@@ -19,5 +19,20 @@
         expect(result).to eq("abc")
       end
     end
+
+    context "when given an ExecutionContext" do
+      let(:aggregator) do
+        described_class.new("ContextAggregator", operation: ->(outputs) { outputs.values.join(", ") })
+      end
+
+      it "passes the context outputs to the operation" do
+        context = MARS::ExecutionContext.new(input: "query")
+        context.record(:step_a, "result_a")
+        context.record(:step_b, "result_b")
+
+        result = aggregator.run(context)
+        expect(result).to eq("result_a, result_b")
+      end
+    end
   end
 end
diff --git a/spec/mars/workflows/parallel_spec.rb b/spec/mars/workflows/parallel_spec.rb
index ab05d25..27142bc 100644
--- a/spec/mars/workflows/parallel_spec.rb
+++ b/spec/mars/workflows/parallel_spec.rb
@@ -1,23 +1,23 @@
 # frozen_string_literal: true
 
 RSpec.describe MARS::Workflows::Parallel do
-  let(:sum_aggregator) { MARS::Aggregator.new("Sum Aggregator", operation: lambda(&:sum)) }
   let(:add_step_class) do
-    Class.new do
-      def initialize(value)
+    Class.new(MARS::Runnable) do
+      def initialize(value:, **kwargs)
+        super(**kwargs)
         @value = value
       end
 
       def run(input)
-        sleep 0.1
         input + @value
       end
     end
   end
 
   let(:multiply_step_class) do
-    Class.new do
-      def initialize(multiplier)
+    Class.new(MARS::Runnable) do
+      def initialize(multiplier:, **kwargs)
+        super(**kwargs)
         @multiplier = multiplier
       end
 
@@ -28,12 +28,10 @@ def run(input)
   end
 
   let(:error_step_class) do
-    Class.new do
-      attr_reader :name
-
-      def initialize(message, name)
+    Class.new(MARS::Runnable) do
+      def initialize(message:, **kwargs)
+        super(**kwargs)
         @message = message
-        @name = name
       end
 
       def run(_input)
@@ -43,44 +41,88 @@ def run(_input)
   end
 
   describe "#run" do
-    it "executes steps in parallel without an aggregator" do
-      add_five = add_step_class.new(5)
-      multiply_three = multiply_step_class.new(3)
-      add_two = add_step_class.new(2)
+    it "executes steps in parallel and returns context with all outputs" do
+      add_five = add_step_class.new(value: 5, name: "add_five")
+      multiply_three = multiply_step_class.new(multiplier: 3, name: "multiply_three")
+      add_two = add_step_class.new(value: 2, name: "add_two")
+
+      aggregator = MARS::Aggregator.new("sum", operation: ->(ctx) { ctx.values.sum })
+      workflow = described_class.new("math_workflow",
+                                     steps: [add_five, multiply_three, add_two],
+                                     aggregator: aggregator)
+
+      # 10+5=15, 10*3=30, 10+2=12 → sum=57
+      result = workflow.run(10)
+      expect(result).to eq(57)
+    end
+
+    it "records each step output in the merged context" do
+      add_five = add_step_class.new(value: 5, name: "add_five")
+      multiply_three = multiply_step_class.new(multiplier: 3, name: "multiply_three")
 
-      workflow = described_class.new("math_workflow", steps: [add_five, multiply_three, add_two])
+      aggregator = MARS::Aggregator.new("pass", operation: ->(outputs) { outputs })
+      workflow = described_class.new("math_workflow",
+                                     steps: [add_five, multiply_three],
+                                     aggregator: aggregator)
 
-      # 10 + 5 = 15, 10 * 3 = 30, 10 + 2 = 12
-      expect(workflow.run(10)).to eq([15, 30, 12])
+      result = workflow.run(10)
+      expect(result).to eq({ add_five: 15, multiply_three: 30 })
     end
 
-    it "executes steps in parallel with a custom aggregator" do
-      add_five = add_step_class.new(5)
-      multiply_three = multiply_step_class.new(3)
-      add_two = add_step_class.new(2)
-      workflow = described_class.new("math_workflow", steps: [add_five, multiply_three, add_two],
-                                                      aggregator: sum_aggregator)
+    it "each branch gets independent current_input" do
+      tracker = Class.new(MARS::Runnable) do
+        def run(input)
+          "saw:#{input}"
+        end
+      end
+
+      step_a = tracker.new(name: "a")
+      step_b = tracker.new(name: "b")
+
+      aggregator = MARS::Aggregator.new("collect", operation: lambda(&:values))
+      workflow = described_class.new("independent", steps: [step_a, step_b], aggregator: aggregator)
 
-      expect(workflow.run(10)).to eq(57)
+      result = workflow.run("original")
+      expect(result).to eq(["saw:original", "saw:original"])
     end
 
-    it "handles single step" do
-      multiply_step = multiply_step_class.new(7)
-      workflow = described_class.new("single_step", steps: [multiply_step])
+    it "shares global_state across branches" do
+      writer = Class.new(MARS::Runnable) do
+        after_run do |ctx, _result, step|
+          ctx.global_state[step.name.to_sym] = true
+        end
+
+        def run(input)
+          input
+        end
+      end
+
+      step_a = writer.new(name: "a")
+      step_b = writer.new(name: "b")
+
+      aggregator = MARS::Aggregator.new("check", operation: ->(outputs) { outputs })
+      workflow = described_class.new("shared_state",
+                                     steps: [step_a, step_b],
+                                     aggregator: aggregator)
+
+      ctx = MARS::ExecutionContext.new(input: "x", global_state: {})
+      workflow.run(ctx)
 
-      expect(workflow.run(6)).to eq([42])
+      expect(ctx.global_state[:a]).to be(true)
+      expect(ctx.global_state[:b]).to be(true)
     end
 
-    it "returns empty array when no steps" do
+    it "returns empty result when no steps" do
       workflow = described_class.new("empty", steps: [])
 
-      expect(workflow.run(42)).to eq([])
+      result = workflow.run(42)
+      expect(result).to eq({})
     end
 
     it "propagates errors from steps" do
-      add_step = add_step_class.new(5)
-      error_step = error_step_class.new("Step failed", "error_step_one")
-      error_step_two = error_step_class.new("Step failed two", "error_step_two")
+      add_step = add_step_class.new(value: 5, name: "add")
+      error_step = error_step_class.new(message: "Step failed", name: "error_step_one")
+      error_step_two = error_step_class.new(message: "Step failed two", name: "error_step_two")
 
       workflow = described_class.new("error_workflow", steps: [add_step, error_step, error_step_two])
 
diff --git a/spec/mars/workflows/sequential_spec.rb b/spec/mars/workflows/sequential_spec.rb
index 45c0783..433ec7c 100644
--- a/spec/mars/workflows/sequential_spec.rb
+++ b/spec/mars/workflows/sequential_spec.rb
@@ -2,8 +2,9 @@
 
 RSpec.describe MARS::Workflows::Sequential do
   let(:add_step_class) do
-    Class.new do
-      def initialize(value)
+    Class.new(MARS::Runnable) do
+      def initialize(value:, **kwargs)
+        super(**kwargs)
         @value = value
       end
 
@@ -14,8 +15,9 @@ def run(input)
   end
 
   let(:multiply_step_class) do
-    Class.new do
-      def initialize(multiplier)
+    Class.new(MARS::Runnable) do
+      def initialize(multiplier:, **kwargs)
+        super(**kwargs)
         @multiplier = multiplier
       end
 
@@ -26,8 +28,9 @@ def run(input)
   end
 
   let(:error_step_class) do
-    Class.new do
-      def initialize(message)
+    Class.new(MARS::Runnable) do
+      def initialize(message:, **kwargs)
+        super(**kwargs)
         @message = message
       end
 
@@ -38,38 +41,129 @@ def run(_input)
   end
 
   describe "#run" do
-    it "executes steps sequentially" do
-      add_five = add_step_class.new(5)
-      multiply_three = multiply_step_class.new(3)
-      add_two = add_step_class.new(2)
+    it "executes steps sequentially and returns context" do
+      add_five = add_step_class.new(value: 5, name: "add_five")
+      multiply_three = multiply_step_class.new(multiplier: 3, name: "multiply_three")
+      add_two = add_step_class.new(value: 2, name: "add_two")
 
       workflow = described_class.new("math_workflow", steps: [add_five, multiply_three, add_two])
+      context = workflow.run(10)
 
       # 10 + 5 = 15, 15 * 3 = 45, 45 + 2 = 47
-      expect(workflow.run(10)).to eq(47)
+      expect(context.current_input).to eq(47)
+    end
+
+    it "records each step output in the context" do
+      add_five = add_step_class.new(value: 5, name: "add_five")
+      multiply_three = multiply_step_class.new(multiplier: 3, name: "multiply_three")
+
+      workflow = described_class.new("math_workflow", steps: [add_five, multiply_three])
+      context = workflow.run(10)
+
+      expect(context[:add_five]).to eq(15)
+      expect(context[:multiply_three]).to eq(45)
     end
 
     it "handles single step" do
-      multiply_step = multiply_step_class.new(7)
+      multiply_step = multiply_step_class.new(multiplier: 7, name: "multiply")
       workflow = described_class.new("single_step", steps: [multiply_step])
 
-      expect(workflow.run(6)).to eq(42)
+      context = workflow.run(6)
+      expect(context.current_input).to eq(42)
     end
 
-    it "returns input unchanged when no steps" do
+    it "returns context with original input when no steps" do
       workflow = described_class.new("empty", steps: [])
 
-      expect(workflow.run(42)).to eq(42)
+      context = workflow.run(42)
+      expect(context.current_input).to eq(42)
     end
 
     it "propagates errors from steps" do
-      add_step = add_step_class.new(5)
-      error_step = error_step_class.new("Step failed")
+      add_step = add_step_class.new(value: 5, name: "add")
+      error_step = error_step_class.new(message: "Step failed", name: "error")
 
       workflow = described_class.new("error_workflow", steps: [add_step, error_step])
 
       expect { workflow.run(10) }.to raise_error(StandardError, "Step failed")
     end
+
+    it "accepts an existing ExecutionContext" do
+      add_step = add_step_class.new(value: 1, name: "add")
+      workflow = described_class.new("ctx_workflow", steps: [add_step])
+
+      ctx = MARS::ExecutionContext.new(input: 100, global_state: { key: "val" })
+      result = workflow.run(ctx)
+
+      expect(result.current_input).to eq(101)
+      expect(result.global_state[:key]).to eq("val")
+    end
+
+    it "runs before and after hooks" do
+      hook_log = []
+      step_class = Class.new(MARS::Runnable) do
+        before_run { |_ctx, step| hook_log << "before:#{step.name}" }
+        after_run { |_ctx, result, step| hook_log << "after:#{step.name}:#{result}" }
+
+        def run(input)
+          input.upcase
+        end
+      end
+
+      # capture hook_log in the closure
+      local_log = hook_log
+      step_class.define_method(:hook_log) { local_log }
+
+      step = step_class.new(name: "upper")
+      workflow = described_class.new("hook_workflow", steps: [step])
+      workflow.run("hello")
+
+      expect(hook_log).to eq(["before:upper", "after:upper:HELLO"])
+    end
+
+    it "applies formatters" do
+      custom_formatter_class = Class.new(MARS::Formatter) do
+        def format_input(context)
+          "prefix:#{context.current_input}"
+        end
+
+        def format_output(output)
+          "#{output}:suffix"
+        end
+      end
+
+      step_class = Class.new(MARS::Runnable) do
+        def run(input)
+          input.upcase
+        end
+      end
+
+      step = step_class.new(name: "fmt_step", formatter: custom_formatter_class.new)
+      workflow = described_class.new("fmt_workflow", steps: [step])
+      context = workflow.run("hello")
+
+      expect(context[:fmt_step]).to eq("PREFIX:HELLO:suffix")
+    end
+  end
+
+  describe ".build" do
+    it "builds a workflow from a block" do
+      step_class = Class.new(MARS::Runnable) do
+        self.step_name = "my_step"
+
+        def run(input)
+          "#{input}!"
+        end
+      end
+      stub_const("ExclaimStep", step_class)
+
+      workflow = described_class.build("built") do
+        step ExclaimStep
+      end
+
+      context = workflow.run("hello")
+      expect(context.current_input).to eq("hello!")
+    end
   end
 
   describe "inheritance" do

From b7c379529ce3a58dde129cba5bcf46bcce633d51 Mon Sep 17 00:00:00 2001
From: Santiago Bartesaghi <sbartesaghi@hotmail.com>
Date: Sun, 1 Mar 2026 13:18:06 -0300
Subject: [PATCH 6/6] Rendering

---
 README.md                                  | 62 ++++++++-------
 examples/complex_llm_workflow/generator.rb | 68 +++++++++--------
 examples/complex_workflow/diagram.html     | 64 ++++++++++++++++
 examples/complex_workflow/diagram.md       | 35 ++++-----
 examples/complex_workflow/generator.rb     | 10 +--
 examples/parallel_workflow/generator.rb    |  6 +-
 examples/simple_workflow/generator.rb      |  6 +-
 lib/mars/rendering/graph.rb                |  1 +
 lib/mars/rendering/graph/base.rb           |  7 ++
 lib/mars/rendering/html.rb                 | 88 ++++++++++++++++++++++
 lib/mars/rendering/mermaid.rb              | 32 +++++++-
 spec/mars/rendering/graph_spec.rb          | 34 +++++++++
 spec/mars/rendering/html_spec.rb           | 73 ++++++++++++++++++
 13 files changed, 392 insertions(+), 94 deletions(-)
 create mode 100644 examples/complex_workflow/diagram.html
 create mode 100644 lib/mars/rendering/html.rb
 create mode 100644 spec/mars/rendering/graph_spec.rb
 create mode 100644 spec/mars/rendering/html_spec.rb

diff --git a/README.md b/README.md
index 0c14456..d0aecca 100644
--- a/README.md
+++ b/README.md
@@ -47,47 +47,53 @@ Here's a simple example to get you started:
 ```ruby
 require 'mars'
 
-# Define agents
-class Agent1 < Mars::Agent
+# Define a RubyLLM agent
+class MyAgent < RubyLLM::Agent
+  model "gpt-4o"
+  instructions "You are a helpful assistant."
 end
 
-class Agent2 < Mars::Agent
+# Wrap it in a MARS step
+class MyStep < MARS::AgentStep
+  agent MyAgent
 end
 
-class Agent3 < Mars::Agent
-end
-
-# Create agents
-agent1 = Agent1.new
-agent2 = Agent2.new
-agent3 = Agent3.new
+# Create steps
+step1 = MyStep.new(name: "step1")
+step2 = MyStep.new(name: "step2")
+step3 = MyStep.new(name: "step3")
 
 # Create a sequential workflow
-workflow = Mars::Workflows::Sequential.new(
+workflow = MARS::Workflows::Sequential.new(
   "My First Workflow",
-  steps: [agent1, agent2, agent3]
+  steps: [step1, step2, step3]
 )
 
 # Run the workflow
-result = workflow.run("Your input here")
+context = workflow.run("Your input here")
+context.current_input  # final output
+context[:step1]        # access any step's output by name
 ```
 
 ## Core Concepts
 
-### Agents
+### Agent Steps
 
-Agents are the basic building blocks of MARS. They represent individual units of work:
+Agent steps are the basic building blocks of MARS. They wrap a `RubyLLM::Agent` subclass for workflow orchestration:
 
 ```ruby
-class CustomAgent < Mars::Agent
-  def system_prompt
-    "You are a helpful assistant"
-  end
+class ResearcherAgent < RubyLLM::Agent
+  model "gpt-4o"
+  instructions "You research topics thoroughly."
+  tools WebSearch
+  schema OutputSchema
 end
 
-agent = CustomAgent.new(
-  options: { model: "gpt-4o" }
-)
+class ResearcherStep < MARS::AgentStep
+  agent ResearcherAgent
+end
+
+step = ResearcherStep.new(name: "researcher")
 ```
 
 ### Sequential Workflows
@@ -95,7 +101,7 @@ agent = CustomAgent.new(
 Execute agents one after another, passing outputs as inputs:
 
 ```ruby
-sequential = Mars::Workflows::Sequential.new(
+sequential = MARS::Workflows::Sequential.new(
   "Sequential Pipeline",
   steps: [agent1, agent2, agent3]
 )
@@ -106,12 +112,12 @@ sequential = Mars::Workflows::Sequential.new(
 Run multiple agents concurrently and aggregate their results:
 
 ```ruby
-aggregator = Mars::Aggregator.new(
+aggregator = MARS::Aggregator.new(
   "Results Aggregator",
   operation: lambda { |results| results.join(", ") }
 )
 
-parallel = Mars::Workflows::Parallel.new(
+parallel = MARS::Workflows::Parallel.new(
   "Parallel Pipeline",
   steps: [agent1, agent2, agent3],
   aggregator: aggregator
@@ -123,7 +129,7 @@ parallel = Mars::Workflows::Parallel.new(
 Create conditional branching in your workflows:
 
 ```ruby
-gate = Mars::Gate.new(
+gate = MARS::Gate.new(
   "Decision Gate",
   condition: ->(input) { input[:score] > 0.5 ? :success : :failure },
   branches: {
@@ -138,7 +144,7 @@ gate = Mars::Gate.new(
 Generate Mermaid diagrams to visualize your workflows:
 
 ```ruby
-diagram = Mars::Rendering::Mermaid.new(workflow).render
+diagram = MARS::Rendering::Mermaid.new(workflow).render
 File.write("workflow_diagram.md", diagram)
 ```
 
@@ -197,7 +203,7 @@ The gem is available as open source under the terms of the [MIT License](https:/
 
 ## Code of Conduct
 
-Everyone interacting in the Mars project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/rootstrap/mars/blob/main/CODE_OF_CONDUCT.md).
+Everyone interacting in the MARS project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/rootstrap/mars/blob/main/CODE_OF_CONDUCT.md).
 
 ## Credits
 
diff --git a/examples/complex_llm_workflow/generator.rb b/examples/complex_llm_workflow/generator.rb
index 4ea2c77..0bfe4b6 100755
--- a/examples/complex_llm_workflow/generator.rb
+++ b/examples/complex_llm_workflow/generator.rb
@@ -37,48 +37,52 @@ def execute(latitude:, longitude:)
   end
 end
 
-# Define LLMs
-class Agent1 < MARS::Agent
-  def system_prompt
-    "You are a helpful assistant that can answer questions.
-     When asked about a country, only answer with its name."
-  end
+# Define RubyLLM agents
+class CountryAgent < RubyLLM::Agent
+  model "gpt-4o"
+  instructions "You are a helpful assistant that can answer questions. " \
+               "When asked about a country, only answer with its name."
 end
 
-class Agent2 < MARS::Agent
-  def system_prompt
-    "You are a helpful assistant that can answer questions and help with tasks.
-     Return information about the typical food of the country."
-  end
+class FoodAgent < RubyLLM::Agent
+  model "gpt-4o"
+  instructions "You are a helpful assistant. Return information about the typical food of the country."
 end
 
-class Agent3 < MARS::Agent
-  def system_prompt
-    "You are a helpful assistant that can answer questions and help with tasks.
-     Return information about the popular sports of the country."
-  end
+class SportsAgent < RubyLLM::Agent
+  model "gpt-4o"
+  instructions "You are a helpful assistant. Return information about the popular sports of the country."
+  schema SportsSchema
+end
 
-  def schema
-    SportsSchema.new
-  end
+class WeatherAgent < RubyLLM::Agent
+  model "gpt-4o"
+  instructions "You are a helpful assistant. Return the current weather of the country's capital."
+  tools Weather
 end
 
-class Agent4 < MARS::Agent
-  def system_prompt
-    "You are a helpful assistant that can answer questions and help with tasks.
-     Return the current weather of the country's capital."
-  end
+# Define MARS steps wrapping RubyLLM agents
+class CountryStep < MARS::AgentStep
+  agent CountryAgent
+end
 
-  def tools
-    [Weather.new]
-  end
+class FoodStep < MARS::AgentStep
+  agent FoodAgent
+end
+
+class SportsStep < MARS::AgentStep
+  agent SportsAgent
+end
+
+class WeatherStep < MARS::AgentStep
+  agent WeatherAgent
 end
 
-# Create the LLMs
-llm1 = Agent1.new(options: { model: "gpt-4o" })
-llm2 = Agent2.new(options: { model: "gpt-4o" })
-llm3 = Agent3.new(options: { model: "gpt-4o" })
-llm4 = Agent4.new(options: { model: "gpt-4o" })
+# Create the steps
+llm1 = CountryStep.new(name: "Country")
+llm2 = FoodStep.new(name: "Food")
+llm3 = SportsStep.new(name: "Sports")
+llm4 = WeatherStep.new(name: "Weather")
 
 parallel_workflow = MARS::Workflows::Parallel.new(
   "Parallel workflow",
diff --git a/examples/complex_workflow/diagram.html b/examples/complex_workflow/diagram.html
new file mode 100644
index 0000000..53b7070
--- /dev/null
+++ b/examples/complex_workflow/diagram.html
@@ -0,0 +1,64 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>Main Pipeline</title>
+  <style>
+    body { font-family: system-ui, sans-serif; margin: 2rem; background: #fafafa; }
+    h1 { color: #333; }
+    #diagram { background: white; padding: 2rem; border-radius: 8px; box-shadow: 0 1px 3px rgba(0,0,0,0.1); }
+  </style>
+</head>
+<body>
+  <h1>Main Pipeline</h1>
+  <div id="diagram"></div>
+  <script type="module">
+    import { renderMermaidSVG } from "https://esm.sh/beautiful-mermaid@1";
+    const diagram = `flowchart LR
+in((In))
+out((Out))
+agent1[agent1]
+gate{Gate}
+agent4[agent4]
+parallel_workflow_aggregator[Parallel workflow Aggregator]
+agent2[agent2]
+agent3[agent3]
+parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator]
+agent5[agent5]
+subgraph main_pipeline["Main Pipeline"]
+  agent1
+  gate
+  parallel_workflow_aggregator
+  subgraph parallel_workflow_2["Parallel workflow 2"]
+    subgraph sequential_workflow["Sequential workflow"]
+      agent4
+      subgraph parallel_workflow["Parallel workflow"]
+        agent2
+        agent3
+      end
+      parallel_workflow_aggregator
+    end
+    agent5
+  end
+  parallel_workflow_2_aggregator
+end
+in --> agent1
+agent1 --> gate
+gate -->|warning| agent4
+gate -->|error| agent2
+gate -->|error| agent3
+gate --> agent4
+gate --> agent5
+agent4 --> agent2
+agent4 --> agent3
+agent2 --> parallel_workflow_aggregator
+parallel_workflow_aggregator --> parallel_workflow_2_aggregator
+agent3 --> parallel_workflow_aggregator
+parallel_workflow_2_aggregator --> out
+agent5 --> parallel_workflow_2_aggregator`;
+    const svg = renderMermaidSVG(diagram);
+    document.getElementById("diagram").innerHTML = svg;
+  </script>
+</body>
+</html>
diff --git a/examples/complex_workflow/diagram.md b/examples/complex_workflow/diagram.md
index a424141..dd8fc85 100644
--- a/examples/complex_workflow/diagram.md
+++ b/examples/complex_workflow/diagram.md
@@ -2,32 +2,29 @@
 flowchart LR
 in((In))
 out((Out))
-agent1[Agent1]
+agent1[agent1]
 gate{Gate}
-agent4[Agent4]
+agent4[agent4]
 parallel_workflow_aggregator[Parallel workflow Aggregator]
-agent2[Agent2]
-agent3[Agent3]
+agent2[agent2]
+agent3[agent3]
 parallel_workflow_2_aggregator[Parallel workflow 2 Aggregator]
-agent5[Agent5]
-subgraph parallel_workflow_2["Parallel workflow 2"]
-  sequential_workflow
-  agent5
-end
-subgraph parallel_workflow["Parallel workflow"]
-  agent2
-  agent3
-end
-subgraph sequential_workflow["Sequential workflow"]
-  agent4
-  parallel_workflow
-  parallel_workflow_aggregator
-end
+agent5[agent5]
 subgraph main_pipeline["Main Pipeline"]
   agent1
   gate
   parallel_workflow_aggregator
-  parallel_workflow_2
+  subgraph parallel_workflow_2["Parallel workflow 2"]
+    subgraph sequential_workflow["Sequential workflow"]
+      agent4
+      subgraph parallel_workflow["Parallel workflow"]
+        agent2
+        agent3
+      end
+      parallel_workflow_aggregator
+    end
+    agent5
+  end
   parallel_workflow_2_aggregator
 end
 in --> agent1
diff --git a/examples/complex_workflow/generator.rb b/examples/complex_workflow/generator.rb
index 62adc6b..d32a2fa 100755
--- a/examples/complex_workflow/generator.rb
+++ b/examples/complex_workflow/generator.rb
@@ -4,19 +4,19 @@
 require_relative "../../lib/mars"
 
 # Define LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
 end
 
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
 end
 
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
 end
 
-class Agent4 < MARS::Agent
+class Agent4 < MARS::AgentStep
 end
 
-class Agent5 < MARS::Agent
+class Agent5 < MARS::AgentStep
 end
 
 # Create the LLMs
diff --git a/examples/parallel_workflow/generator.rb b/examples/parallel_workflow/generator.rb
index 795c2c5..66378bd 100755
--- a/examples/parallel_workflow/generator.rb
+++ b/examples/parallel_workflow/generator.rb
@@ -4,13 +4,13 @@
 require_relative "../../lib/mars"
 
 # Define the LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
 end
 
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
 end
 
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
 end
 
 # Create the LLMs
diff --git a/examples/simple_workflow/generator.rb b/examples/simple_workflow/generator.rb
index 4dd01c4..b1a0351 100755
--- a/examples/simple_workflow/generator.rb
+++ b/examples/simple_workflow/generator.rb
@@ -4,13 +4,13 @@
 require_relative "../../lib/mars"
 
 # Define the LLMs
-class Agent1 < MARS::Agent
+class Agent1 < MARS::AgentStep
 end
 
-class Agent2 < MARS::Agent
+class Agent2 < MARS::AgentStep
 end
 
-class Agent3 < MARS::Agent
+class Agent3 < MARS::AgentStep
 end
 
 # Create the LLMs
diff --git a/lib/mars/rendering/graph.rb b/lib/mars/rendering/graph.rb
index b2ab350..c453a81 100644
--- a/lib/mars/rendering/graph.rb
+++ b/lib/mars/rendering/graph.rb
@@ -4,6 +4,7 @@ module MARS
   module Rendering
     module Graph
       def self.include_extensions
+        MARS::Runnable.include(Base)
         MARS::AgentStep.include(AgentStep)
         MARS::Gate.include(Gate)
         MARS::Workflows::Sequential.include(SequentialWorkflow)
diff --git a/lib/mars/rendering/graph/base.rb b/lib/mars/rendering/graph/base.rb
index 2e70f54..d785470 100644
--- a/lib/mars/rendering/graph/base.rb
+++ b/lib/mars/rendering/graph/base.rb
@@ -21,6 +21,13 @@ def node_id
           @node_id ||= sanitize(name)
         end
 
+        def to_graph(builder, parent_id: nil, value: nil)
+          builder.add_node(node_id, name, Node::STEP)
+          builder.add_edge(parent_id, node_id, value)
+
+          [node_id]
+        end
+
         private
 
         def sanitize(name)
diff --git a/lib/mars/rendering/html.rb b/lib/mars/rendering/html.rb
new file mode 100644
index 0000000..09c1297
--- /dev/null
+++ b/lib/mars/rendering/html.rb
@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module MARS
+  module Rendering
+    class Html
+      BEAUTIFUL_MERMAID_URL = "https://esm.sh/beautiful-mermaid@1"
+
+      attr_reader :obj
+
+      def initialize(obj)
+        @obj = obj
+      end
+
+      def render(options = {})
+        mermaid = Mermaid.new(obj)
+        diagram = mermaid.graph_mermaid.join("\n")
+        direction = options.fetch(:direction, "LR")
+        title = options.fetch(:title, obj.name)
+        theme = options.fetch(:theme, {})
+
+        build_html(title, direction, diagram, theme)
+      end
+
+      def write(path, options = {})
+        File.write(path, render(options))
+      end
+
+      private
+
+      def build_html(title, direction, diagram, theme)
+        <<~HTML
+          <!DOCTYPE html>
+          <html lang="en">
+          <head>
+            <meta charset="UTF-8">
+            <meta name="viewport" content="width=device-width, initial-scale=1.0">
+            <title>#{escape(title)}</title>
+            #{head_style}
+          </head>
+          <body>
+            <h1>#{escape(title)}</h1>
+            <div id="diagram"></div>
+            #{render_script(direction, diagram, theme)}
+          </body>
+          </html>
+        HTML
+      end
+
+      def head_style
+        <<~STYLE.chomp
+          <style>
+              body { font-family: system-ui, sans-serif; margin: 2rem; background: #fafafa; }
+              h1 { color: #333; }
+              #diagram { background: white; padding: 2rem; border-radius: 8px; box-shadow: 0 1px 3px rgba(0,0,0,0.1); }
+            </style>
+        STYLE
+      end
+
+      def render_script(direction, diagram, theme)
+        theme_opts = theme_options(theme)
+
+        <<~SCRIPT.chomp
+          <script type="module">
+              import { renderMermaidSVG } from "#{BEAUTIFUL_MERMAID_URL}";
+              const diagram = `flowchart #{direction}\n#{escape_js(diagram)}`;
+              const svg = renderMermaidSVG(diagram#{theme_opts});
+              document.getElementById("diagram").innerHTML = svg;
+            </script>
+        SCRIPT
+      end
+
+      def theme_options(theme)
+        return "" if theme.empty?
+
+        pairs = theme.map { |k, v| "#{k}: '#{escape_js(v.to_s)}'" }
+        ", { #{pairs.join(", ")} }"
+      end
+
+      def escape(text)
+        text.to_s.gsub("&", "&amp;").gsub("<", "&lt;").gsub(">", "&gt;").gsub('"', "&quot;")
+      end
+
+      def escape_js(text)
+        text.to_s.gsub("\\", "\\\\\\\\").gsub("`", "\\`").gsub("$", "\\$")
+      end
+    end
+  end
+end
diff --git a/lib/mars/rendering/mermaid.rb b/lib/mars/rendering/mermaid.rb
index b9f2ac8..3f1914a 100644
--- a/lib/mars/rendering/mermaid.rb
+++ b/lib/mars/rendering/mermaid.rb
@@ -31,10 +31,8 @@ def nodes_mermaid
       end
 
       def subgraphs_mermaid
-        subgraphs.values.reverse.map do |subgraph|
-          node_names = subgraph.nodes
-          "subgraph #{subgraph.id}[\"#{subgraph.name}\"]\n  #{node_names.join("\n  ")}\nend"
-        end
+        root_ids = subgraphs.keys - nested_subgraph_ids
+        root_ids.map { |id| render_subgraph(id) }
       end
 
       def edges_mermaid
@@ -66,6 +64,32 @@ def edge_value(value)
 
         "|#{value}|"
       end
+
+      private
+
+      def nested_subgraph_ids
+        ids = []
+        subgraphs.each_value do |sg|
+          sg.nodes.each { |n| ids << n if subgraphs.key?(n) }
+        end
+        ids
+      end
+
+      def render_subgraph(id, indent = "")
+        sg = subgraphs[id]
+        lines = ["#{indent}subgraph #{sg.id}[\"#{sg.name}\"]"]
+        sg.nodes.each { |node_id| lines << render_subgraph_node(node_id, indent) }
+        lines << "#{indent}end"
+        lines.join("\n")
+      end
+
+      def render_subgraph_node(node_id, indent)
+        if subgraphs.key?(node_id)
+          render_subgraph(node_id, "#{indent}  ")
+        else
+          "#{indent}  #{node_id}"
+        end
+      end
     end
   end
 end
diff --git a/spec/mars/rendering/graph_spec.rb b/spec/mars/rendering/graph_spec.rb
new file mode 100644
index 0000000..17b1bfe
--- /dev/null
+++ b/spec/mars/rendering/graph_spec.rb
@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+RSpec.describe MARS::Rendering::Graph do
+  describe "default Runnable rendering" do
+    it "renders any Runnable subclass as a box node" do
+      step_class = Class.new(MARS::Runnable) do
+        def run(input)
+          input
+        end
+      end
+
+      step = step_class.new(name: "custom_step")
+      graph, nodes, _subgraphs = step.build_graph
+
+      expect(nodes).to have_key("custom_step")
+      expect(nodes["custom_step"].type).to eq(MARS::Rendering::Graph::Node::STEP)
+      expect(graph).to have_key("in")
+    end
+  end
+
+  describe "AgentStep rendering" do
+    it "renders as a step node" do
+      mock_agent = Class.new
+      step_class = Class.new(MARS::AgentStep) do
+        agent mock_agent
+      end
+
+      step = step_class.new(name: "my_agent")
+      _graph, nodes, _subgraphs = step.build_graph
+
+      expect(nodes["my_agent"].type).to eq(MARS::Rendering::Graph::Node::STEP)
+    end
+  end
+end
diff --git a/spec/mars/rendering/html_spec.rb b/spec/mars/rendering/html_spec.rb
new file mode 100644
index 0000000..d09b9a9
--- /dev/null
+++ b/spec/mars/rendering/html_spec.rb
@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "tmpdir"
+
+RSpec.describe MARS::Rendering::Html do
+  let(:step_class) do
+    Class.new(MARS::Runnable) do
+      def run(input)
+        input
+      end
+    end
+  end
+
+  let(:workflow) do
+    step_a = step_class.new(name: "step_a")
+    step_b = step_class.new(name: "step_b")
+    MARS::Workflows::Sequential.new("TestPipeline", steps: [step_a, step_b])
+  end
+
+  describe "#render" do
+    it "returns a self-contained HTML string" do
+      html = described_class.new(workflow).render
+
+      expect(html).to include("<!DOCTYPE html>")
+      expect(html).to include("beautiful-mermaid")
+      expect(html).to include("renderMermaidSVG")
+      expect(html).to include("TestPipeline")
+      expect(html).to include("flowchart LR")
+    end
+
+    it "accepts a custom direction" do
+      html = described_class.new(workflow).render(direction: "TD")
+
+      expect(html).to include("flowchart TD")
+    end
+
+    it "accepts a custom title" do
+      html = described_class.new(workflow).render(title: "My Workflow")
+
+      expect(html).to include("<title>My Workflow</title>")
+    end
+
+    it "accepts theme options" do
+      html = described_class.new(workflow).render(theme: { bg: "#1a1b26", fg: "#a9b1d6" })
+
+      expect(html).to include("bg: '#1a1b26'")
+      expect(html).to include("fg: '#a9b1d6'")
+    end
+
+    it "escapes HTML in the title" do
+      html = described_class.new(workflow).render(title: "<script>alert('xss')</script>")
+
+      expect(html).not_to include("<script>alert")
+      expect(html).to include("&lt;script&gt;")
+    end
+  end
+
+  describe "#write" do
+    it "writes the HTML to a file" do
+      path = File.join(Dir.tmpdir, "mars_test_#{SecureRandom.hex(4)}.html")
+
+      begin
+        described_class.new(workflow).write(path)
+
+        content = File.read(path)
+        expect(content).to include("<!DOCTYPE html>")
+        expect(content).to include("beautiful-mermaid")
+      ensure
+        FileUtils.rm_f(path)
+      end
+    end
+  end
+end