Skip to content

fix: Explicitly setting the otel parent spans in agents, llm flow and function calls #998

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Mar 11, 2026
Merged

fix: Explicitly setting the otel parent spans in agents, llm flow and function calls #998

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.errorprone.annotations.DoNotCall;
import com.google.genai.types.Content;
import io.opentelemetry.context.Context;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
Expand Down Expand Up @@ -311,6 +312,7 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
private Flowable<Event> run(
InvocationContext parentContext,
Function<InvocationContext, Flowable<Event>> runImplementation) {
Context parentSpanContext = Context.current();
return Flowable.defer(
() -> {
InvocationContext invocationContext = createInvocationContext(parentContext);
Expand Down Expand Up @@ -339,8 +341,12 @@ private Flowable<Event> run(
})
.switchIfEmpty(mainAndAfterEvents)
.compose(
Tracing.traceAgent(
"invoke_agent " + name(), name(), description(), invocationContext));
Tracing.<Event>trace("invoke_agent " + name())
.setParent(parentSpanContext)
.configure(
span ->
Tracing.traceAgentInvocation(
span, name(), description(), invocationContext)));
});
}

Expand Down
27 changes: 18 additions & 9 deletions core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,10 @@ protected Flowable<Event> postprocess(
* callbacks. Callbacks should not rely on its ID if they create their own separate events.
*/
private Flowable<LlmResponse> callLlm(
InvocationContext context, LlmRequest llmRequest, Event eventForCallbackUsage) {
Context spanContext,
InvocationContext context,
LlmRequest llmRequest,
Event eventForCallbackUsage) {
LlmAgent agent = (LlmAgent) context.agent();

LlmRequest.Builder llmRequestBuilder = llmRequest.toBuilder();
Expand Down Expand Up @@ -200,7 +203,7 @@ private Flowable<LlmResponse> callLlm(
span.setStatus(StatusCode.ERROR, error.getMessage());
span.recordException(error);
})
.compose(Tracing.trace("call_llm"))
.compose(Tracing.<LlmResponse>trace("call_llm").setParent(spanContext))
.concatMap(
llmResp ->
handleAfterModelCallback(context, llmResp, eventForCallbackUsage)
Expand Down Expand Up @@ -319,7 +322,7 @@ private Single<LlmResponse> handleAfterModelCallback(
* @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations.
* @throws IllegalStateException if a transfer agent is specified but not found.
*/
private Flowable<Event> runOneStep(InvocationContext context) {
private Flowable<Event> runOneStep(Context spanContext, InvocationContext context) {
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());

return Flowable.defer(
Expand Down Expand Up @@ -351,7 +354,11 @@ private Flowable<Event> runOneStep(InvocationContext context) {
.build();
mutableEventTemplate.setTimestamp(0L);

return callLlm(context, llmRequestAfterPreprocess, mutableEventTemplate)
return callLlm(
spanContext,
context,
llmRequestAfterPreprocess,
mutableEventTemplate)
.concatMap(
llmResponse -> {
try (Scope postScope = currentContext.makeCurrent()) {
Expand Down Expand Up @@ -403,11 +410,12 @@ private Flowable<Event> runOneStep(InvocationContext context) {
*/
@Override
public Flowable<Event> run(InvocationContext invocationContext) {
return run(invocationContext, 0);
return run(Context.current(), invocationContext, 0);
}

private Flowable<Event> run(InvocationContext invocationContext, int stepsCompleted) {
Flowable<Event> currentStepEvents = runOneStep(invocationContext).cache();
private Flowable<Event> run(
Context spanContext, InvocationContext invocationContext, int stepsCompleted) {
Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
if (stepsCompleted + 1 >= maxSteps) {
logger.debug("Ending flow execution because max steps reached.");
return currentStepEvents;
Expand All @@ -427,7 +435,7 @@ private Flowable<Event> run(InvocationContext invocationContext, int stepsComple
return Flowable.empty();
} else {
logger.debug("Continuing to next step of the flow.");
return run(invocationContext, stepsCompleted + 1);
return run(spanContext, invocationContext, stepsCompleted + 1);
}
}));
}
Expand All @@ -444,6 +452,7 @@ private Flowable<Event> run(InvocationContext invocationContext, int stepsComple
public Flowable<Event> runLive(InvocationContext invocationContext) {
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
Flowable<Event> preprocessEvents = preprocess(invocationContext, llmRequestRef);
Context spanContext = Context.current();

return preprocessEvents.concatWith(
Flowable.defer(
Expand Down Expand Up @@ -481,7 +490,7 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
eventIdForSendData,
llmRequestAfterPreprocess.contents());
})
.compose(Tracing.trace("send_data"));
.compose(Tracing.<Event>trace("send_data").setParent(spanContext));

Flowable<LiveRequest> liveRequests =
invocationContext
Expand Down
10 changes: 6 additions & 4 deletions core/src/main/java/com/google/adk/flows/llmflows/Functions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ public static Maybe<Event> handleFunctionCalls(
if (events.size() > 1) {
return Maybe.just(mergedEvent)
.doOnSuccess(event -> Tracing.traceToolResponse(event.id(), event))
.compose(Tracing.trace("tool_response", parentContext));
.compose(Tracing.<Event>trace("tool_response").setParent(parentContext));
}
return Maybe.just(mergedEvent);
});
Expand Down Expand Up @@ -432,8 +432,8 @@ private static Maybe<Event> postProcessFunctionResult(
toolContext,
invocationContext))
.compose(
Tracing.trace(
"tool_response [" + tool.name() + "]", parentContext))
Tracing.<Event>trace("tool_response [" + tool.name() + "]")
.setParent(parentContext))
.doOnSuccess(event -> Tracing.traceToolResponse(event.id(), event));
});
}
Expand Down Expand Up @@ -593,7 +593,9 @@ private static Maybe<Map<String, Object>> callTool(
Tracing.traceToolCall(
tool.name(), tool.description(), tool.getClass().getSimpleName(), args))
.doOnError(t -> Span.current().recordException(t))
.compose(Tracing.trace("tool_call [" + tool.name() + "]", parentContext))
.compose(
Tracing.<Map<String, Object>>trace("tool_call [" + tool.name() + "]")
.setParent(parentContext))
.onErrorResumeNext(
e ->
Maybe.error(
Expand Down
13 changes: 0 additions & 13 deletions core/src/main/java/com/google/adk/telemetry/Tracing.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,19 +426,6 @@ public static <T> TracerProvider<T> trace(String spanName) {
return new TracerProvider<>(spanName);
}

/**
* Returns a transformer that traces the execution of an RxJava stream with an explicit parent
* context.
*
* @param spanName The name of the span to create.
* @param parentContext The explicit parent context for the span.
* @param <T> The type of the stream.
* @return A TracerProvider that can be used with .compose().
*/
public static <T> TracerProvider<T> trace(String spanName, Context parentContext) {
return new TracerProvider<T>(spanName).setParent(parentContext);
}

/**
* Returns a transformer that traces an agent invocation.
*
Expand Down