ailine/backend/app/agent/stream_context.py

"""
流式上下文，用于在 LangGraph 节点和 agent_service 之间传递 token 队列
清晰的 API，更易用！
"""
import contextvars
import asyncio
from typing import Optional, Any

# 上下文变量：存储每个请求专属的 token 队列
stream_queue_ctx: contextvars.ContextVar[Optional[asyncio.Queue]] = contextvars.ContextVar(
    "stream_queue", default=None
)


def set_stream_queue(queue: asyncio.Queue) -> None:
    """设置当前请求的队列"""
    stream_queue_ctx.set(queue)


def get_stream_queue() -> Optional[asyncio.Queue]:
    """获取当前请求的队列"""
    return stream_queue_ctx.get()