fix(smolagents): Getting Multiple Traces While Streaming Code Agent

[ialisaleh]

This PR fixes an issue where CodeAgent.run(stream=True) in smolagents would generate separate traces for each step instead of grouping them under a single parent trace. The problem stemmed from Python generators not preserving OpenTelemetry context across yield statements. To address this, a new wrapped_generator() was introduced that ensures the tracing context is reattached during streaming and the final output is collected for proper span annotation. This change preserves existing behavior for non-streaming runs and ensures consistent, complete tracing data regardless of the stream flag.

Closes #1322

[cursor]

Bug: Telemetry Span Status Overwritten Incorrectly

The finally blocks in both streaming and non-streaming execution paths unconditionally set the OpenTelemetry span status to OK. This overwrites any ERROR status previously set by an except block when an exception occurs, causing failed operations to be incorrectly reported as successful in telemetry.

python/instrumentation/openinference-instrumentation-smolagents/src/openinference/instrumentation/smolagents/_wrappers.py#L152-L222

openinference/python/instrumentation/openinference-instrumentation-smolagents/src/openinference/instrumentation/smolagents/_wrappers.py

Lines 152 to 222 in 8848820

	def wrapped_generator() -> Generator[str, None, None]:
	try:
	# Collect chunks for final output
	for chunk in agent_output:
	output_chunks.append(str(chunk))
	yield chunk
	except Exception as e:
	span.record_exception(e)
	span.set_status(trace_api.StatusCode.ERROR)
	raise
	finally:
	# Set output value from the last observation
	steps = getattr(agent.monitor, "steps", [])
	history = getattr(agent.monitor, "history", [])
	if steps:
	observation = getattr(steps[-1], "observations", None)
	if observation:
	span.set_attribute(OUTPUT_VALUE, observation)
	elif history:
	observation = getattr(history[-1], "observations", None)
	if observation:
	span.set_attribute(OUTPUT_VALUE, observation)
	elif output_chunks:
	span.set_attribute(OUTPUT_VALUE, "".join(output_chunks))
	# Record token usage metadata
	span.set_attribute(
	LLM_TOKEN_COUNT_PROMPT, agent.monitor.total_input_token_count
	)
	span.set_attribute(
	LLM_TOKEN_COUNT_COMPLETION, agent.monitor.total_output_token_count
	)
	span.set_attribute(
	LLM_TOKEN_COUNT_TOTAL,
	agent.monitor.total_input_token_count
	+ agent.monitor.total_output_token_count,
	)
	span.set_status(trace_api.StatusCode.OK)
	span.end()
	context_api.detach(token)
	return wrapped_generator()
	# Handle non-streaming (normal) run
	else:
	try:
	# Set output value from the agent output
	span.set_attribute(OUTPUT_VALUE, str(agent_output))
	# Record token usage metadata
	span.set_attribute(LLM_TOKEN_COUNT_PROMPT, agent.monitor.total_input_token_count)
	span.set_attribute(
	LLM_TOKEN_COUNT_COMPLETION, agent.monitor.total_output_token_count
	)
	span.set_attribute(
	LLM_TOKEN_COUNT_TOTAL,
	agent.monitor.total_input_token_count + agent.monitor.total_output_token_count,
	)
	return agent_output
	except Exception as e:
	span.record_exception(e)
	span.set_status(trace_api.StatusCode.ERROR)
	raise
	finally:
	span.set_status(trace_api.StatusCode.OK)
	span.end()
	context_api.detach(token)

Fix in Cursor • Fix in Web

---

Bug: Unhandled Exceptions Cause Resource Leaks

A resource leak occurs if the wrapped function raises an exception. The context_api token is attached and a span is started, but the initial exception handler records the error and re-raises without detaching the context token or ending the span, leading to leaked resources.

python/instrumentation/openinference-instrumentation-smolagents/src/openinference/instrumentation/smolagents/_wrappers.py#L139-L145

openinference/python/instrumentation/openinference-instrumentation-smolagents/src/openinference/instrumentation/smolagents/_wrappers.py

Lines 139 to 145 in 8848820

	try:
	agent_output = wrapped(*args, **kwargs)
	except Exception as e:
	span.record_exception(e)
	span.set_status(trace_api.StatusCode.ERROR)
	raise

Fix in Cursor • Fix in Web

---

Bug: Error Status Overwritten in Finally Block

In the streaming generator's finally block, the span status is unconditionally set to OK. This overwrites any ERROR status previously set by the except block when an exception occurs during generator iteration, leading to a loss of error information on the span.

python/instrumentation/openinference-instrumentation-smolagents/src/openinference/instrumentation/smolagents/_wrappers.py#L191-L194

openinference/python/instrumentation/openinference-instrumentation-smolagents/src/openinference/instrumentation/smolagents/_wrappers.py

Lines 191 to 194 in 8848820

	span.set_status(trace_api.StatusCode.OK)
	span.end()
	context_api.detach(token)

Fix in Cursor • Fix in Web

---

Was this report helpful? Give feedback by reacting with 👍 or 👎

[axiomofjoy]

This test looks like it is being xfailed and doesn't seem to test the emitted traces at all. Can we nix altogether and add a more minimal test for the generator?