From 9a080763d83c319f539d1bacac4595d13b299e7e Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Thu, 19 Mar 2026 07:11:45 -0700 Subject: [PATCH] feat: fixing context propagation for agent transfers PiperOrigin-RevId: 886159283 --- .../adk/flows/llmflows/BaseLlmFlow.java | 16 +- .../java/com/google/adk/telemetry/README.md | 156 ++++++++++ .../adk/telemetry/ContextPropagationTest.java | 269 +++++++++++++----- 3 files changed, 368 insertions(+), 73 deletions(-) create mode 100644 core/src/main/java/com/google/adk/telemetry/README.md diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index e00cf0cbf..8fabc978d 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -430,7 +430,12 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex "Agent not found: " + agentToTransfer))); } return postProcessedEvents.concatWith( - Flowable.defer(() -> nextAgent.get().runAsync(context))); + Flowable.defer( + () -> { + try (Scope s = spanContext.makeCurrent()) { + return nextAgent.get().runAsync(context); + } + })); } return postProcessedEvents; }); @@ -488,6 +493,8 @@ private Flowable run( public Flowable runLive(InvocationContext invocationContext) { AtomicReference llmRequestRef = new AtomicReference<>(LlmRequest.builder().build()); Flowable preprocessEvents = preprocess(invocationContext, llmRequestRef); + // Capture agent context at assembly time to use as parent for agent transfer at subscription + // time. See Flowable.defer() usages below. Context spanContext = Context.current(); return preprocessEvents.concatWith( @@ -608,7 +615,12 @@ public void onError(Throwable e) { "Agent not found: " + event.actions().transferToAgent().get()); } Flowable nextAgentEvents = - nextAgent.get().runLive(invocationContext); + Flowable.defer( + () -> { + try (Scope s = spanContext.makeCurrent()) { + return nextAgent.get().runLive(invocationContext); + } + }); events = Flowable.concat(events, nextAgentEvents); } return events; diff --git a/core/src/main/java/com/google/adk/telemetry/README.md b/core/src/main/java/com/google/adk/telemetry/README.md new file mode 100644 index 000000000..8665b3352 --- /dev/null +++ b/core/src/main/java/com/google/adk/telemetry/README.md @@ -0,0 +1,156 @@ +# ADK Telemetry and Tracing + +This package contains classes for capturing and reporting telemetry data within +the ADK, primarily for tracing agent execution leveraging OpenTelemetry. + +## Overview + +The `Tracing` utility class provides methods to trace various aspects of an +agent's execution, including: + +* Agent invocations +* LLM requests and responses +* Tool calls and responses + +These traces can be exported and visualized in telemetry backends like Google +Cloud Trace or Zipkin, or viewed through the ADK Dev Server UI, providing +observability into agent behavior. + +## How Tracing is Used + +Tracing is deeply integrated into the ADK's RxJava-based asynchronous workflows. + +### Agent Invocations + +Every agent's `runAsync` or `runLive` execution is wrapped in a span named +`invoke_agent `. The top-level agent invocation initiated by +`Runner.runAsync` or `Runner.runLive` is captured in a span named `invocation`. +Agent-specific metadata like name and description are added as span attributes, +following OpenTelemetry semantic conventions (e.g., `gen_ai.agent.name`). + +### LLM Calls + +Calls to Large Language Models (LLMs) are traced within a `call_llm` span. The +`traceCallLlm` method attaches detailed attributes to this span, including: + +* The LLM request (excluding large data like images) and response. +* Model name (`gen_ai.request.model`). +* Token usage (`gen_ai.usage.input_tokens`, `gen_ai.usage.output_tokens`). +* Configuration parameters (`gen_ai.request.top_p`, + `gen_ai.request.max_tokens`). +* Response finish reason (`gen_ai.response.finish_reasons`). + +### Tool Calls and Responses + +Tool executions triggered by the LLM are traced using `tool_call []` +and `tool_response []` spans. + +* `traceToolCall` records tool arguments in the + `gcp.vertex.agent.tool_call_args` attribute. +* `traceToolResponse` records tool output in the + `gcp.vertex.agent.tool_response` attribute. +* If multiple tools are called in parallel, a single `tool_response` span may + be created for the merged result. + +### Context Propagation + +ADK is built on RxJava and heavily uses asynchronous processing, which means +that work is often handed off between different threads. For tracing to work +correctly in such an environment, it's crucial that the active span's context +is propagated across these thread boundaries. If context is not propagated, +new spans may be orphaned or attached to the wrong parent, making traces +difficult to interpret. + +OpenTelemetry stores the currently active span in a thread-local variable. +When an asynchronous operation switches threads, this thread-local context is +lost. To solve this, ADK's `Tracing` class provides functionality to capture +the context on one thread and restore it on another when an asynchronous +operation resumes. This ensures that spans created on different threads are +correctly parented under the same trace. + +The primary mechanism for this is the `Tracing.withContext(context)` method, +which returns an RxJava transformer. When applied to an RxJava stream via +`.compose()`, this transformer ensures that the provided `Context` (containing +the parent span) is re-activated before any `onNext`, `onError`, `onComplete`, +or `onSuccess` signals are propagated downstream. It achieves this by wrapping +the downstream observer with a `TracingObserver`, which uses +`context.makeCurrent()` in a try-with-resources block around each callback, +guaranteeing that the correct span is active when downstream operators execute, +regardless of the thread. + +### RxJava Integration + +ADK integrates OpenTelemetry with RxJava streams to simplify span creation and +ensure context propagation: + +* **Span Creation**: The `Tracing.trace(spanName)` method returns an RxJava + transformer that can be applied to a `Flowable`, `Single`, `Maybe`, or + `Completable` using `.compose()`. This transformer wraps the stream's + execution in a new OpenTelemetry span. +* **Context Propagation**: The `Tracing.withContext(context)` transformer is + used with `.compose()` to ensure that the correct OpenTelemetry `Context` + (and thus the correct parent span) is active when stream operators or + subscriptions are executed, even across thread boundaries. + +## Trace Hierarchy Example + +A typical agent interaction might produce a trace hierarchy like the following: + +``` +invocation +└── invoke_agent my_agent + ├── call_llm + │ ├── tool_call [search_flights] + │ └── tool_response [search_flights] + └── call_llm +``` + +This shows: + +1. The overall `invocation` started by the `Runner`. +2. The invocation of `my_agent`. +3. The first `call_llm` made by `my_agent`. +4. A `tool_call` to `search_flights` and its corresponding `tool_response`. +5. A second `call_llm` made by `my_agent` to generate the final user response. + +### Nested Agents + +ADK supports nested agents, where one agent invokes another. If an agent has +sub-agents, it can transfer control to one of them using the built-in +`transfer_to_agent` tool. When `AgentA` calls `transfer_to_agent` to transfer +control to `AgentB`, the `invoke_agent AgentB` span will appear as a child of +the `invoke_agent AgentA` span, like so: + +``` +invocation +└── invoke_agent AgentA + ├── call_llm + │ ├── tool_call [transfer_to_agent] + │ └── tool_response [transfer_to_agent] + └── invoke_agent AgentB + ├── call_llm + └── ... +``` + +This structure allows you to see how `AgentA` delegated work to `AgentB`. + +## Span Creation References + +The following classes are the primary places where spans are created: + +* **`com.google.adk.runner.Runner`**: Initiates the top-level `invocation` + span for `runAsync` and `runLive`. +* **`com.google.adk.agents.BaseAgent`**: Creates the `invoke_agent + ` span for each agent execution. +* **`com.google.adk.flows.llmflows.BaseLlmFlow`**: Creates the `call_llm` span + when the LLM is invoked. +* **`com.google.adk.flows.llmflows.Functions`**: Creates `tool_call [...]` and + `tool_response [...]` spans when handling tool calls and responses. + +## Configuration + +**ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS**: This environment variable controls +whether LLM request/response content and tool arguments/responses are captured +in span attributes. It defaults to `true`. Set to `false` to exclude potentially +large or sensitive data from traces, in which case a `{}` JSON object will be +recorded instead. diff --git a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java index e5795d61f..1ee018848 100644 --- a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java +++ b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java @@ -16,6 +16,7 @@ package com.google.adk.telemetry; +import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -32,10 +33,15 @@ import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; import com.google.adk.sessions.SessionKey; +import com.google.adk.testing.TestLlm; +import com.google.adk.testing.TestUtils; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.FunctionResponse; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponseUsageMetadata; @@ -54,6 +60,7 @@ import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.Optional; @@ -123,10 +130,7 @@ public void testToolCallSpanLinksToParent() { parentSpanData.getSpanContext().getTraceId(), toolCallSpanData.getSpanContext().getTraceId()); - assertEquals( - "Tool call's parent should be the parent span", - parentSpanData.getSpanContext().getSpanId(), - toolCallSpanData.getParentSpanContext().getSpanId()); + assertParent(parentSpanData, toolCallSpanData); } @Test @@ -146,7 +150,7 @@ public void testToolCallWithoutParentCreatesRootSpan() { // Then: Should create root span (backward compatible) List spans = openTelemetryRule.getSpans(); - assertEquals("Should have exactly 1 span", 1, spans.size()); + assertThat(spans).hasSize(1); SpanData toolCallSpanData = spans.get(0); assertFalse( @@ -193,7 +197,7 @@ public void testNestedSpanHierarchy() { List spans = openTelemetryRule.getSpans(); // The 4 spans are: "parent", "invocation", "tool_call [testTool]", and "tool_response // [testTool]". - assertEquals("Should have 4 spans in the hierarchy", 4, spans.size()); + assertThat(spans).hasSize(4); SpanData parentSpanData = findSpanByName("parent"); String parentTraceId = parentSpanData.getSpanContext().getTraceId(); @@ -210,22 +214,13 @@ public void testNestedSpanHierarchy() { SpanData toolResponseSpanData = findSpanByName("tool_response [testTool]"); // invocation should be child of parent - assertEquals( - "Invocation should be child of parent", - parentSpanData.getSpanContext().getSpanId(), - invocationSpanData.getParentSpanContext().getSpanId()); + assertParent(parentSpanData, invocationSpanData); // tool_call should be child of invocation - assertEquals( - "Tool call should be child of invocation", - invocationSpanData.getSpanContext().getSpanId(), - toolCallSpanData.getParentSpanContext().getSpanId()); + assertParent(invocationSpanData, toolCallSpanData); // tool_response should be child of tool_call - assertEquals( - "Tool response should be child of tool call", - toolCallSpanData.getSpanContext().getSpanId(), - toolResponseSpanData.getParentSpanContext().getSpanId()); + assertParent(toolCallSpanData, toolResponseSpanData); } @Test @@ -253,7 +248,6 @@ public void testMultipleSpansInParallel() { // Verify all tool calls link to same parent SpanData parentSpanData = findSpanByName("parent"); String parentTraceId = parentSpanData.getSpanContext().getTraceId(); - String parentSpanId = parentSpanData.getSpanContext().getSpanId(); // All tool calls should have same trace ID and parent span ID List toolCallSpans = @@ -261,7 +255,7 @@ public void testMultipleSpansInParallel() { .filter(s -> s.getName().startsWith("tool_call")) .toList(); - assertEquals("Should have 3 tool call spans", 3, toolCallSpans.size()); + assertThat(toolCallSpans).hasSize(3); toolCallSpans.forEach( span -> { @@ -269,10 +263,7 @@ public void testMultipleSpansInParallel() { "Tool call should have same trace ID as parent", parentTraceId, span.getSpanContext().getTraceId()); - assertEquals( - "Tool call should have parent as parent span", - parentSpanId, - span.getParentSpanContext().getSpanId()); + assertParent(parentSpanData, span); }); } @@ -298,10 +289,7 @@ public void testInvokeAgentSpanLinksToInvocation() { SpanData invocationSpanData = findSpanByName("invocation"); SpanData invokeAgentSpanData = findSpanByName("invoke_agent test-agent"); - assertEquals( - "Agent run should be child of invocation", - invocationSpanData.getSpanContext().getSpanId(), - invokeAgentSpanData.getParentSpanContext().getSpanId()); + assertParent(invocationSpanData, invokeAgentSpanData); } @Test @@ -323,15 +311,12 @@ public void testCallLlmSpanLinksToAgentRun() { } List spans = openTelemetryRule.getSpans(); - assertEquals("Should have 2 spans", 2, spans.size()); + assertThat(spans).hasSize(2); SpanData invokeAgentSpanData = findSpanByName("invoke_agent test-agent"); SpanData callLlmSpanData = findSpanByName("call_llm"); - assertEquals( - "Call LLM should be child of agent run", - invokeAgentSpanData.getSpanContext().getSpanId(), - callLlmSpanData.getParentSpanContext().getSpanId()); + assertParent(invokeAgentSpanData, callLlmSpanData); } @Test @@ -349,10 +334,7 @@ public void testSpanCreatedWithinParentScopeIsCorrectlyParented() { SpanData parentSpanData = findSpanByName("invocation"); SpanData agentSpanData = findSpanByName("invoke_agent"); - assertEquals( - "Agent span should be a child of the invocation span", - parentSpanData.getSpanContext().getSpanId(), - agentSpanData.getParentSpanContext().getSpanId()); + assertParent(parentSpanData, agentSpanData); } @Test @@ -380,9 +362,7 @@ public void testTraceFlowable() throws InterruptedException { SpanData parentSpanData = findSpanByName("parent"); SpanData flowableSpanData = findSpanByName("flowable"); - assertEquals( - parentSpanData.getSpanContext().getSpanId(), - flowableSpanData.getParentSpanContext().getSpanId()); + assertParent(parentSpanData, flowableSpanData); assertTrue(flowableSpanData.hasEnded()); } @@ -469,9 +449,7 @@ public void testTraceTransformer() throws InterruptedException { SpanData parentSpanData = findSpanByName("parent"); SpanData transformerSpanData = findSpanByName("transformer"); - assertEquals( - parentSpanData.getSpanContext().getSpanId(), - transformerSpanData.getParentSpanContext().getSpanId()); + assertParent(parentSpanData, transformerSpanData); assertTrue(transformerSpanData.hasEnded()); } @@ -485,7 +463,7 @@ public void testTraceAgentInvocation() { span.end(); } List spans = openTelemetryRule.getSpans(); - assertEquals(1, spans.size()); + assertThat(spans).hasSize(1); SpanData spanData = spans.get(0); Attributes attrs = spanData.getAttributes(); assertEquals("invoke_agent", attrs.get(AttributeKey.stringKey("gen_ai.operation.name"))); @@ -504,7 +482,7 @@ public void testTraceToolCall() { span.end(); } List spans = openTelemetryRule.getSpans(); - assertEquals(1, spans.size()); + assertThat(spans).hasSize(1); SpanData spanData = spans.get(0); Attributes attrs = spanData.getAttributes(); assertEquals("execute_tool", attrs.get(AttributeKey.stringKey("gen_ai.operation.name"))); @@ -541,7 +519,7 @@ public void testTraceToolResponse() { span.end(); } List spans = openTelemetryRule.getSpans(); - assertEquals(1, spans.size()); + assertThat(spans).hasSize(1); SpanData spanData = spans.get(0); Attributes attrs = spanData.getAttributes(); assertEquals("execute_tool", attrs.get(AttributeKey.stringKey("gen_ai.operation.name"))); @@ -578,7 +556,7 @@ public void testTraceCallLlm() { span.end(); } List spans = openTelemetryRule.getSpans(); - assertEquals(1, spans.size()); + assertThat(spans).hasSize(1); SpanData spanData = spans.get(0); Attributes attrs = spanData.getAttributes(); assertEquals("gcp.vertex.agent", attrs.get(AttributeKey.stringKey("gen_ai.system"))); @@ -606,12 +584,12 @@ public void testTraceSendData() { Tracing.traceSendData( buildInvocationContext(), "event-1", - ImmutableList.of(Content.builder().role("user").parts(Part.fromText("hello")).build())); + ImmutableList.of(Content.fromParts(Part.fromText("hello")))); } finally { span.end(); } List spans = openTelemetryRule.getSpans(); - assertEquals(1, spans.size()); + assertThat(spans).hasSize(1); SpanData spanData = spans.get(0); Attributes attrs = spanData.getAttributes(); assertEquals( @@ -653,37 +631,23 @@ public void baseAgentRunAsync_propagatesContext() throws InterruptedException { } SpanData parent = findSpanByName("parent"); SpanData agentSpan = findSpanByName("invoke_agent test-agent"); - assertEquals(parent.getSpanContext().getSpanId(), agentSpan.getParentSpanContext().getSpanId()); + assertParent(parent, agentSpan); } @Test public void runnerRunAsync_propagatesContext() throws InterruptedException { BaseAgent agent = new TestAgent(); - Runner runner = Runner.builder().agent(agent).appName("test-app").build(); Span parentSpan = tracer.spanBuilder("parent").startSpan(); try (Scope s = parentSpan.makeCurrent()) { - Session session = - runner - .sessionService() - .createSession(new SessionKey("test-app", "test-user", "test-session")) - .blockingGet(); - Content newMessage = Content.fromParts(Part.fromText("hi")); - RunConfig runConfig = RunConfig.builder().build(); - runner - .runAsync(session.userId(), session.id(), newMessage, runConfig, null) - .test() - .await() - .assertComplete(); + runAgent(agent); } finally { parentSpan.end(); } SpanData parent = findSpanByName("parent"); SpanData invocation = findSpanByName("invocation"); SpanData agentSpan = findSpanByName("invoke_agent test-agent"); - assertEquals( - parent.getSpanContext().getSpanId(), invocation.getParentSpanContext().getSpanId()); - assertEquals( - invocation.getSpanContext().getSpanId(), agentSpan.getParentSpanContext().getSpanId()); + assertParent(parent, invocation); + assertParent(invocation, agentSpan); } @Test @@ -713,10 +677,173 @@ public void runnerRunLive_propagatesContext() throws InterruptedException { SpanData parent = findSpanByName("parent"); SpanData invocation = findSpanByName("invocation"); SpanData agentSpan = findSpanByName("invoke_agent test-agent"); - assertEquals( - parent.getSpanContext().getSpanId(), invocation.getParentSpanContext().getSpanId()); - assertEquals( - invocation.getSpanContext().getSpanId(), agentSpan.getParentSpanContext().getSpanId()); + assertParent(parent, invocation); + assertParent(invocation, agentSpan); + } + + @Test + public void testAgentWithToolCallTraceHierarchy() throws InterruptedException { + // This test verifies the trace hierarchy created when an agent calls an LLM, + // which then invokes a tool. The expected hierarchy is: + // invocation + // └── invoke_agent test_agent + // ├── call_llm + // │ ├── tool_call [search_flights] + // │ └── tool_response [search_flights] + // └── call_llm + + SearchFlightsTool searchFlightsTool = new SearchFlightsTool(); + + TestLlm testLlm = + TestUtils.createTestLlm( + TestUtils.createLlmResponse( + Content.builder() + .role("model") + .parts( + Part.fromFunctionCall( + searchFlightsTool.name(), ImmutableMap.of("destination", "SFO"))) + .build()), + TestUtils.createLlmResponse(Content.fromParts(Part.fromText("done")))); + + LlmAgent agentWithTool = + LlmAgent.builder() + .name("test_agent") + .description("description") + .model(testLlm) + .tools(ImmutableList.of(searchFlightsTool)) + .build(); + + runAgent(agentWithTool); + + SpanData invocation = findSpanByName("invocation"); + SpanData invokeAgent = findSpanByName("invoke_agent test_agent"); + SpanData toolCall = findSpanByName("tool_call [search_flights]"); + SpanData toolResponse = findSpanByName("tool_response [search_flights]"); + List callLlmSpans = + openTelemetryRule.getSpans().stream() + .filter(s -> s.getName().equals("call_llm")) + .sorted(Comparator.comparing(SpanData::getStartEpochNanos)) + .toList(); + assertThat(callLlmSpans).hasSize(2); + SpanData callLlm1 = callLlmSpans.get(0); + SpanData callLlm2 = callLlmSpans.get(1); + + // Assert hierarchy: + // invocation + // └── invoke_agent test_agent + assertParent(invocation, invokeAgent); + // ├── call_llm 1 + assertParent(invokeAgent, callLlm1); + // │ ├── tool_call [search_flights] + assertParent(callLlm1, toolCall); + // │ └── tool_response [search_flights] + assertParent(callLlm1, toolResponse); + // └── call_llm 2 + assertParent(invokeAgent, callLlm2); + } + + @Test + public void testNestedAgentTraceHierarchy() throws InterruptedException { + // This test verifies the trace hierarchy created when AgentA transfers to AgentB. + // The expected hierarchy is: + // invocation + // └── invoke_agent AgentA + // ├── call_llm + // │ ├── tool_call [transfer_to_agent] + // │ └── tool_response [transfer_to_agent] + // └── invoke_agent AgentB + // └── call_llm + TestLlm llm = + TestUtils.createTestLlm( + TestUtils.createLlmResponse( + Content.builder() + .role("model") + .parts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "AgentB"))) + .build()), + TestUtils.createLlmResponse(Content.fromParts(Part.fromText("agent b response")))); + LlmAgent agentB = LlmAgent.builder().name("AgentB").description("Agent B").model(llm).build(); + + LlmAgent agentA = + LlmAgent.builder() + .name("AgentA") + .description("Agent A") + .model(llm) + .subAgents(ImmutableList.of(agentB)) + .build(); + + runAgent(agentA); + + SpanData invocation = findSpanByName("invocation"); + SpanData agentASpan = findSpanByName("invoke_agent AgentA"); + SpanData toolCall = findSpanByName("tool_call [transfer_to_agent]"); + SpanData agentBSpan = findSpanByName("invoke_agent AgentB"); + SpanData toolResponse = findSpanByName("tool_response [transfer_to_agent]"); + + List callLlmSpans = + openTelemetryRule.getSpans().stream() + .filter(s -> s.getName().equals("call_llm")) + .sorted(Comparator.comparing(SpanData::getStartEpochNanos)) + .toList(); + assertThat(callLlmSpans).hasSize(2); + + SpanData agentACallLlm1 = callLlmSpans.get(0); + SpanData agentBCallLlm = callLlmSpans.get(1); + + assertParent(invocation, agentASpan); + assertParent(agentASpan, agentACallLlm1); + assertParent(agentACallLlm1, toolCall); + assertParent(agentACallLlm1, toolResponse); + assertParent(agentASpan, agentBSpan); + assertParent(agentBSpan, agentBCallLlm); + } + + private void runAgent(BaseAgent agent) throws InterruptedException { + Runner runner = Runner.builder().agent(agent).appName("test-app").build(); + Session session = + runner + .sessionService() + .createSession(new SessionKey("test-app", "test-user", "test-session")) + .blockingGet(); + Content newMessage = Content.fromParts(Part.fromText("hi")); + RunConfig runConfig = RunConfig.builder().build(); + runner + .runAsync(session.sessionKey(), newMessage, runConfig, null) + .test() + .await() + .assertComplete(); + } + + /** Tool for testing. */ + public static class SearchFlightsTool extends BaseTool { + public SearchFlightsTool() { + super("search_flights", "Search for flights tool"); + } + + @Override + public Single> runAsync(Map args, ToolContext context) { + return Single.just(ImmutableMap.of("result", args)); + } + + @Override + public Optional declaration() { + return Optional.of( + FunctionDeclaration.builder() + .name("search_flights") + .description("Search for flights tool") + .build()); + } + } + + /** + * Asserts that the parent span is the parent of the child span. + * + * @param parent The parent span. + * @param child The child span. + */ + private void assertParent(SpanData parent, SpanData child) { + assertEquals(parent.getSpanContext().getSpanId(), child.getParentSpanContext().getSpanId()); } /**