BUG: FinalAnswerStep callbacks not executed

[victorlearned]

Problem

When using CodeAgent.run() with step_callbacks registered for FinalAnswerStep, the callbacks are never executed. This occurs because FinalAnswerStep is never passed to _finalize_step(), which is responsible for invoking step callbacks.

The issue is in two places:

1. _finalize_step() method signature only accepts ActionStep | PlanningStep, not FinalAnswerStep
2. In _run_stream(), FinalAnswerStep is created and yielded directly without calling _finalize_step() on it

Steps to reproduce

1. Create a CodeAgent with a callback registered for FinalAnswerStep
2. Call agent.run() (non-streaming mode)
3. Observe that the callback is never executed

from smolagents import CodeAgent, FinalAnswerStep, ActionStep

def test_final_answer_callback(step: FinalAnswerStep, agent):
    """Callback that should be triggered when FinalAnswerStep is created."""
    print(f"Callback triggered for FinalAnswerStep: {step.output}")
    step.output = "MODIFIED BY CALLBACK"

# Create agent with step_callbacks
agent = CodeAgent(
    tools=[],
    model=model,  # your model instance
    step_callbacks={FinalAnswerStep: test_final_answer_callback},
)

# Run the agent
result = agent.run("What is 2+2?")

# The callback should have modified the output, but it doesn't
# result will be the original output, not "MODIFIED BY CALLBACK"

Actual behavior and error logs

The callback registered for FinalAnswerStep is never executed. No errors are raised, but the callback simply doesn't run.

# No errors, but callback is silent
# The output remains unchanged by the callback

Expected behavior

When FinalAnswerStep is created in _run_stream(), it should be passed to _finalize_step() before being yielded, which would trigger any callbacks registered for FinalAnswerStep. The callback should execute and be able to modify the step if needed.

Environment:

Please complete the following information:

• OS: macOS (darwin 24.6.0)
• Python version: 3.11
• Package version: (output of pip show smolagents)

# Please run: pip show smolagents

Additional context (optional)

Root Cause Analysis:

1. Location 1: _finalize_step() method (line 572 in agents.py)

   • Current signature: def _finalize_step(self, memory_step: ActionStep | PlanningStep)
   • Issue: FinalAnswerStep is not included in the type hint, though the method would work if called

2. Location 2: _run_stream() method (line 563 in agents.py)

   • Current code:

     yield FinalAnswerStep(handle_agent_output_types(final_answer))

   • Issue: FinalAnswerStep is created and yielded directly without calling _finalize_step()

3. Comparison with other step types:

   • ActionStep is finalized at line 555: self._finalize_step(action_step)
   • PlanningStep is finalized at line 520: self._finalize_step(planning_step)
   • FinalAnswerStep is never finalized

Proposed Fix:

1. Update _finalize_step() signature to accept FinalAnswerStep:

   def _finalize_step(self, memory_step: ActionStep | PlanningStep | FinalAnswerStep):

2. Finalize FinalAnswerStep before yielding in _run_stream():

   final_answer_step = FinalAnswerStep(handle_agent_output_types(final_answer))
   self._finalize_step(final_answer_step)
   yield final_answer_step

Code References:

• _finalize_step(): .venv/lib/python3.11/site-packages/smolagents/agents.py:572-574
• _run_stream() FinalAnswerStep creation: .venv/lib/python3.11/site-packages/smolagents/agents.py:563
• run() method that calls _run_stream(): .venv/lib/python3.11/site-packages/smolagents/agents.py:456-492

---

Checklist

• I have searched the existing issues and have not found a similar bug report.
• I have provided a minimal, reproducible example.
• I have provided the full traceback of the error.
• I have provided my environment details.
• I am willing to work on this issue and submit a pull request. (optional)

[suryabdev]

Hi @victorlearned, There is a field called final_answer_checks that runs on termination. I think that should solve your concern

smolagents/src/smolagents/agents.py

Line 588 in 2ae00fb

if self.final_answer_checks:

https://huggingface.co/docs/smolagents/en/guided_tour#customizing-agent-termination-conditions
Code example from the docs

from smolagents import CodeAgent, InferenceClientModel

# Define a custom final answer check function
def is_integer(final_answer: str, agent_memory=None) -> bool:
    """Return True if final_answer is an integer."""
    try:
        int(final_answer)
        return True
    except ValueError:
        return False

# Initialize agent with custom final answer check
agent = CodeAgent(
    tools=[],
    model=InferenceClientModel(),
    final_answer_checks=[is_integer]
)

agent.run("Calculate the least common multiple of 3 and 7")

[victorlearned]

Hi there — just to give more context on how I’m using this today:

I already use final_answer_checks for situations where the agent returns structured JSON or other typed outputs. It’s helpful for making sure the final answer matches the schema or type I expect, so I’m familiar with the intended workflow.

The only limitation I’m running into is that this API is designed specifically for lightweight validation. Its purpose is essentially to answer the question: “Does the final output match the shape or type I expect?”—and return True or False based on that.

Because it’s built for that narrow role, it doesn’t seem like the right place for more involved logic, such as post-processing the final answer, updating or enriching the agent’s memory, performing any bookkeeping, or applying more complex termination-time behavior. While it’s technically possible to encode more information into the final answer, that doesn’t change the fact that the callback itself is meant as a simple validator rather than a post-processing hook.

Since the memory-step callbacks already exist for working with agent state, it feels more natural for termination to be treated as another lifecycle event that those callbacks can handle. That would provide a cleaner path for cases where the final step needs to update memory or apply structured post-processing, without having to push that work through a boolean validation function.

Happy to elaborate if helpful.

[suryabdev]

step_callbacks (`list[Callable]` | `dict[Type[MemoryStep], Callable | list[Callable]]`, *optional*): Callbacks that will be called at each step.
...
final_answer_checks (`list[Callable]`, *optional*): List of validation functions to run before accepting a final answer.
        Each function should:
        - Take the final answer, the agent's memory, and the agent itself as arguments.
        - Return a boolean indicating whether the final answer is valid.

step_callbacks take the memory_step and the agent as arguments.
final_answer_checks take the final_answer, agent memory and the agent itself as arguments. Plus it has to return a bool indicating whether a validation passed

this API is designed specifically for lightweight validation

Since the memory and agent are passed to the final_answer_check, I'd say it can do everything the step_callback (Updating memory etc) can plus it has the additional responsibility of being a validation function.

Clarifying,

1. If you have to post-process the final answer, Can you do it after the agent returns the value? Is there some use case I'm missing?
2. Using the agent object and memory in the final_answer_callback, Is there some use case you can't implement?

[victorlearned]

Thanks for the follow-up — here’s some clarification based on my use case and how the API behaves.

On doing post-processing after the agent returns

I can technically do that, but it doesn’t work well for my setup. This agent responds to customers over SMS, and the pipeline already includes model latency, Twilio latency, and carrier latency. Adding another post-processing step after run() adds extra delay.

I also want all logs and traces for a single exchange to stay inside one agent execution. If I bolt post-processing on after run(), I end up with two separate execution paths to correlate in logging/tracing, instead of one clean, continuous trace. Keeping everything inside the run avoids that.

On putting all logic inside final_answer_checks

It’s true that final_answer_checks receives the final answer, memory, and the agent object, so it can reach everything it needs. The issue isn’t whether it’s possible — it’s that this hook has a very specific contract: return a boolean indicating whether the answer is valid.

If I use this hook to also handle transformation, logging, enrichment, or other termination-time behavior, I’m stapling several different responsibilities into one callback that was only intended for validation. If I want to keep the logic modular, I’d have to build my own mini dispatcher inside that hook just to chain multiple functions together, even though the agent already provides a clean callback pipeline (step_callbacks) for exactly that.

Why I’m asking for FinalAnswerStep to go through _finalize_step()

The final memory step currently just stores the raw output of the last ActionStep. That’s exactly where I need to hook in and shape or enrich the final output before it becomes the committed answer.

If FinalAnswerStep flowed through the same callback mechanism as other steps, I could break the termination logic into small, focused functions and reuse the existing callback structure instead of packing everything into a validator or adding a second processing layer after the run.

So while it’s technically possible to do everything inside final_answer_checks or after the agent completes, both approaches introduce unnecessary coupling or latency. Running FinalAnswerStep through the existing callback pipeline would keep the design consistent, modular, and compatible with the tracing and latency requirements of my use case.

[suryabdev]

@victorlearned thanks for clarifications, I agree it would make sense not to overload the validation function
I raised a PR #1881, Will check what the others think