fix: Fix: google-adk tool/LLM spans locked at OK on ADK 2.x failures

[satyadevai]

Closes #3415

[caroger]

[P1] Do not finalize OK before ADK callbacks complete

Reproduced on this head (639b8d9) with google-adk==2.6.3: a final response followed by after_model_callback raising RuntimeError("after-model failed") exports the call_llm span as OK with two exception events.

Suggested change:

-            if not llm_response.partial:
-                # This is the final chunk for this LLM turn, so no further
-                # trace_call_llm calls will land on this span.
-                span.set_status(StatusCode.OK)

Success can remain UNSET; this lets the owning ADK span record a later callback failure as ERROR. A regression test can use the existing fixtures:

class _FinalResponseLlm(BaseLlm):
    async def generate_content_async(
        self, llm_request: LlmRequest, stream: bool = False
    ) -> AsyncGenerator[LlmResponse, None]:
        yield LlmResponse(
            content=types.Content(role="model", parts=[types.Part(text="done")])
        )


def _fail_after_model(_ctx: Any, _response: LlmResponse) -> None:
    raise RuntimeError("after-model failed")


async def test_call_llm_span_errors_when_after_model_callback_raises(
    instrument: Any, in_memory_span_exporter: InMemorySpanExporter
) -> None:
    agent = Agent(
        name="agent",
        model=_FinalResponseLlm(model="fake-model"),
        after_model_callback=_fail_after_model,
    )
    runner = InMemoryRunner(agent=agent, app_name="app")
    await runner.session_service.create_session(
        app_name="app", user_id="user", session_id="session"
    )

    with pytest.raises(RuntimeError, match="after-model failed"):
        async for _ in runner.run_async(
            user_id="user",
            session_id="session",
            new_message=types.Content(
                role="user", parts=[types.Part(text="hello")]
            ),
        ):
            pass

    [span] = [
        span
        for span in in_memory_span_exporter.get_finished_spans()
        if span.name == "call_llm"
    ]
    assert span.status.status_code == trace_api.StatusCode.ERROR
    assert any(event.name == "exception" for event in span.events)

[satyadevai]

@caroger As described in the issue, we set the call_llm span to OK once the LLM call has completed with a full, successful response. If an after_model_callback subsequently fails, that exception belongs to the agent flow rather than to the LLM call itself, so marking the completed LLM span as OK is appropriate.

At the moment, we do not have a way to prevent ADK from attaching those later callback exception events to the call_llm span. Because of that, an exception event may appear on a span whose status is still OK, even though the exception was not caused by the LLM call.

For streaming responses, if the LLM response is only partial and the call does not complete successfully, we set the call_llm span to ERROR because the LLM call itself did not complete.

The parent/agent spans are marked as ERROR for the callback failure, which is the appropriate level to represent this failure, while keeping the successfully completed call_llm span as OK.

[satyadevai]

If we remove setting OK when llm_response.partial is false, we effectively revert this part of the code to the state before this ticket. The purpose of this change is to explicitly mark the LLM span as successful once the LLM has returned a complete, successful response.

A later after_model_callback failure is part of the agent flow, not a failure of the already-completed LLM call. So leaving the LLM span as UNSET just to allow a later callback exception to mark it as ERROR would bring back the original issue and incorrectly couple the LLM call status to downstream agent processing.

[caroger]

Requesting changes for one remaining correctness gap in the LLM status fix.

The tool-failure and partial-stream failure paths are fixed, and the Google ADK 2.6.3 lane passes locally (Ruff, mypy, and 34 tests). However, _TraceCallLlm still assigns OK on a final response before ADK invokes after_model_callback. A callback exception therefore produces an OK call_llm span containing exception events—the same class of status-locking bug this PR is intended to eliminate.

The exact-head reproduction and suggested minimal code/test changes are in the inline comment.