Fix premature end of langchain agent span

[qnnn]

LangChain sends the same run_id for agent_action, agent_finish, and chain_end during an agent interaction.

Therefore, only the chain_end callback is used to end the agent span to avoid it being prematurely ended.

---

Important

Removes on_agent_action and on_agent_finish methods from CallbackHandler.py to prevent premature span termination, relying solely on on_chain_end to end spans.

• Behavior:
  • Removes on_agent_action and on_agent_finish methods from CallbackHandler.py.
  • on_chain_end now solely ends the agent span, preventing premature span termination.
• Imports:
  • Removes unused imports AgentAction and AgentFinish from CallbackHandler.py.

^{This description was created by for 2781062. You can customize this summary. It will automatically update as commits are pushed.}

Disclaimer: Experimental PR review

Greptile Summary

This PR fixes a critical issue in the LangChain callback handler where agent spans were being prematurely terminated. The root cause was that LangChain reuses the same run_id across multiple callback events (agent_action, agent_finish, and chain_end) during a single agent interaction, causing spans to be ended multiple times.

The solution removes two callback methods (on_agent_action and on_agent_finish) from the CallbackHandler class in langfuse/langchain/CallbackHandler.py. These methods were calling .end() on spans when they received their respective events, but since all three events share the same run_id, this resulted in spans being closed before the agent interaction was actually complete.

By eliminating these intermediate callbacks and relying solely on on_chain_end to handle span termination, the change ensures that agent spans remain active throughout the entire agent execution lifecycle and are only closed when the agent's work is truly finished. The PR also removes the now-unused imports for AgentAction and AgentFinish from langchain.schema.agent.

This change integrates well with the existing codebase architecture, as the callback handler already has a robust span management system through the self.runs dictionary that tracks active spans by their run_id. The fix aligns with LangChain's callback model where chain_end represents the definitive completion of an agent's execution cycle.

Confidence score: 4/5

• This PR addresses a well-defined bug with a targeted solution that has minimal risk of side effects
• Score reflects the focused nature of the change and clear understanding of LangChain's callback architecture
• Pay close attention to testing with agent-based workflows to ensure the fix works as expected in production scenarios

[CLAassistant]

All committers have signed the CLA.

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Bot Settings | Greptile}

[hassiebp]

@qnnn - thanks for your contribution! Could you please open an issue with / share a minimum repro example of a chain that is broken with the current SDK and will be fixed with this PR?

[qnnn]

@qnnn - thanks for your contribution! Could you please open an issue with / share a minimum repro example of a chain that is broken with the current SDK and will be fixed with this PR?

Sure! I’ve opened issue #8718 and added a minimal reproducible unit test there. Please take a look. This PR removes the span update operations in on_agent_action and on_agent_finish callbacks, so that the input, output, and errors are set in on_chain_start, on_chain_end, and on_chain_error respectively, preventing them from being overwritten.

[hassiebp]

Thanks for your contribution, added in #1312