fix: keep a strong reference to the websocket stream task in autogen-studio

[kratos0718]

Why are these changes needed?

The run websocket handler in autogen-studio (web/routes/ws.py) starts the stream with a bare asyncio.create_task(...) and discards the returned task:

# Start the stream in a separate task
asyncio.create_task(ws_manager.start_stream(run_id, task, team_config))

Per the asyncio docs, the event loop only keeps a weak reference to a task. A task with no other reference can be garbage-collected before it finishes — here that means a user's streaming run can be silently dropped mid-flight under GC pressure, with no error surfaced.

Fix

Keep a strong reference to the task in a module-level set and remove it on completion via add_done_callback (the pattern recommended in the asyncio docs):

stream_task = asyncio.create_task(ws_manager.start_stream(run_id, task, team_config))
_background_tasks.add(stream_task)
stream_task.add_done_callback(_background_tasks.discard)

No behavior change beyond ensuring the stream task is retained until it completes. Single file.

Related issue number

N/A — small self-contained reliability fix.

Checks

• I've made sure the change compiles (python -m py_compile).
• The change is minimal and focused on one issue.

[avinashkamat48-design]

Keeping a strong reference fixes the GC risk, but this still drops task failures silently. If ws_manager.start_stream(...) raises, the done callback only discards the task from _background_tasks; nobody retrieves task.exception(), logs it, or notifies the websocket/client. Could the callback inspect/log failures so background stream errors do not become invisible?

[kratos0718]

Good catch @avinashkamat48-design — you're right, the strong reference alone still let failures pass silently. I've updated the done-callback to log any non-cancelled task exception via logger.error(..., exc_info=...), so a failing start_stream(...) now surfaces in the logs instead of being dropped:

def _on_stream_done(t: asyncio.Task, run_id: int = run_id) -> None:
    _background_tasks.discard(t)
    if not t.cancelled() and (exc := t.exception()) is not None:
        logger.error(f"Stream task for run {run_id} failed: {exc}", exc_info=exc)

stream_task.add_done_callback(_on_stream_done)

Happy to route it to the client over the socket instead of (or in addition to) logging if you'd prefer that. Thanks for the review!

[kratos0718]

@microsoft-github-policy-service agree