Fix the intermediate steps span logic to work better with nested coroutines and tasks

[mdemoret-nv]

Description

This PR updates the logic to handle if the END event was called from a Task. This situation breaks the previous logic because tasks create a copy of the context preventing contextvar.set() from working.

By Submitting this PR I confirm:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

[Copilot]

Pull Request Overview

This PR improves the asynchronous handling of alert triage workflows by fixing issues related to context propagation with nested coroutines and tasks. It converts the assistant function to an async method with proper await of LLM calls, adds type annotations for clearer typing, and integrates profiling by tracking execution of the alert process function.

Comments suppressed due to low confidence (2)

examples/alert_triage_agent/src/aiq_alert_triage_agent/register.py:92

• Converting 'ata_assistant' to an async function means that all its call sites must now await its result; please verify that its consumers have been updated accordingly.

async def ata_assistant(state: MessagesState):

examples/alert_triage_agent/src/aiq_alert_triage_agent/register.py:119

• The addition of the '@track_function()' decorator aids profiling, but please ensure that any performance overhead it introduces has been evaluated for production usage.

@track_function()

[AnuradhaKaruppiah]

/merge

[Copilot]

Pull Request Overview

Fixes the intermediate span-tracking logic to correctly restore context when END events come from nested coroutines or separate tasks, and updates tests to validate via a subscribed output list using the new IntermediateStepType enum.

• Manager now records both previous and active stacks in OpenStep and restores or mutates them depending on coroutine vs task context.
• Tests subscribe to manager outputs in output_steps and use IntermediateStepType instead of the old event_state.
• Example workflow updated to type-annotate the LLM, switch assistant helper to async, use ainvoke, and add a track_function decorator.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
tests/aiq/builder/test_intermediate_step_manager.py	Replaced context var stream with output_steps fixture and updated payload to use IntermediateStepType.
src/aiq/builder/intermediate_step_manager.py	Added prev_stack/active_stack to OpenStep, modified START/END handling to fully restore span stack across tasks.
examples/alert_triage_agent/src/aiq_alert_triage_agent/register.py	Annotated llm as BaseChatModel, made assistant function async, switched to ainvoke, and added @track_function.

Comments suppressed due to low confidence (1)

src/aiq/builder/intermediate_step_manager.py:65

• The manager still branches on payload.event_state, but tests now only set payload.event_type. Update these checks to use payload.event_type == IntermediateStepType.LLM_START (and similar) or ensure event_state is populated from event_type.

if (payload.event_state == IntermediateStepState.START):

[Copilot]

This loop can raise an IndexError if parent_step_id is not found. Add a guard like while curr_stack and curr_stack[-1] != parent_step_id: to avoid popping past an empty list.

Suggested change

	while (curr_stack[-1] != parent_step_id):
	curr_stack.pop()
	pop_count += 1
	if (pop_count != 1):
	while curr_stack and curr_stack[-1] != parent_step_id:
	curr_stack.pop()
	pop_count += 1
	if not curr_stack:
	logger.error(
	"Parent step id %s not found in the current stack. "
	"This indicates a potential inconsistency.",
	parent_step_id)
	elif pop_count != 1:

[Copilot]

[nitpick] The output_steps fixture actually collects IntermediateStep objects (not IntermediateStepPayload), so the annotation can be misleading. Consider changing it to list[IntermediateStep] for clarity.

Suggested change

	def mgr_fixture(ctx_state: AIQContextState, output_steps: list[IntermediateStepPayload]):
	def mgr_fixture(ctx_state: AIQContextState, output_steps: list[IntermediateStep]):