"""Streaming agent-hierarchical path — never block Zed on a silent await.""" from __future__ import annotations import asyncio import json import logging import time from typing import Any import httpx from fastapi.responses import StreamingResponse log = logging.getLogger("ai-router.agent_stream") def sse_chunk( *, cid: str, model: str, delta: dict[str, Any], finish_reason: str | None = None, ) -> bytes: payload = { "id": cid, "object": "chat.completion.chunk", "created": int(time.time()), "model": model, "choices": [ { "index": 0, "delta": delta, "finish_reason": finish_reason, } ], } return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n".encode("utf-8") def format_progress_block(lines: list[str]) -> str: from progress_ui import format_progress_block as _fmt return _fmt(lines) async def stream_agent_plan_then_act( *, runner: Any, last_user: str, quality_mode: str | None, session_id: str, messages: list[dict[str, Any]], body: dict[str, Any], meta: dict[str, Any], decision: Any, prompt_hash: str, litellm_url: str, litellm_key: str, cfg: dict[str, Any], find_pending_plan: Any, plan_payload_from_meta: Any, inject_plan_context: Any, pick_agent_executor_model: Any, prepare_agent_executor_forward: Any, completion_to_sse_chunks: Any, synthetic_first_tool_completion: Any, meta_header: Any, orchestrator: Any, requests_metric: Any, executor_fallback_completion: Any = None, ) -> StreamingResponse: """Immediately SSE-stream progress; plan in background; never silent-hang.""" cid = f"hier-agent-{int(time.time() * 1000)}" queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue() async def on_progress(msg: str) -> None: await queue.put(("progress", msg)) async def work() -> None: try: hier = await runner.run( last_user, quality_mode=quality_mode, session_id=session_id, messages=messages, on_progress=on_progress, agent_mode=True, ) await queue.put(("done", hier)) except Exception as exc: # noqa: BLE001 log.exception("agent plan failed: %s", exc) await queue.put(("error", exc)) async def gen(): # First bytes ASAP — Zed must leave spinner for content yield sse_chunk( cid=cid, model="hierarchical", delta={ "role": "assistant", "content": "**Agent**\n· Планирую задачу\n", }, ) task = asyncio.create_task(work()) hier = None failed: Exception | None = None progress_lines: list[str] = ["Планирую задачу"] while True: try: kind, payload = await asyncio.wait_for(queue.get(), timeout=12.0) except asyncio.TimeoutError: # SSE comment keepalive (proxies / Zed idle timeout) yield b": keepalive\n\n" if task.done() and queue.empty(): # drain outcome if any if not task.cancelled(): try: _ = task.result() except Exception as exc: # noqa: BLE001 failed = exc break continue if kind == "progress": from progress_ui import humanize_line, stream_step line = humanize_line(str(payload)) if not line: continue progress_lines.append(line) yield sse_chunk( cid=cid, model="hierarchical", delta={"content": stream_step(line)}, ) continue if kind == "error": failed = payload # type: ignore[assignment] break hier = payload break if not task.done(): await task if failed is not None: err = (str(failed) or type(failed).__name__)[:200] meta["hierarchical_error"] = err yield sse_chunk( cid=cid, model="hierarchical", delta={ "content": ( f"\n---\nОшибка плана (роутер, не LiteLLM UI): `{err}`.\n" "Повтори запрос или напиши проще. Спиннер без шагов = баг " "роутера — мы должны стримить progress; если снова тишина, " "проверь деплой.\n" ) }, ) yield sse_chunk( cid=cid, model="hierarchical", delta={}, finish_reason="stop" ) yield b"data: [DONE]\n\n" return if hier is None: yield sse_chunk( cid=cid, model="hierarchical", delta={"content": "\n---\nПустой ответ плана.\n"}, ) yield sse_chunk( cid=cid, model="hierarchical", delta={}, finish_reason="stop" ) yield b"data: [DONE]\n\n" return hier_meta = {**meta, **hier.meta} if hier.meta.get("agent_execute"): plan = plan_payload_from_meta(hier.meta) or find_pending_plan(messages) orchestrator.after_request( session_id, prompt_hash=prompt_hash, success=True, escalate=False ) from path_resolve import ( apply_deterministic_path_index, path_resolve_needed, pick_path_resolve_model, ) if cfg.get("path_resolve_enabled", True): plan, unresolved = apply_deterministic_path_index(plan) if unresolved: log.info("path_index miss (stream) → %s", unresolved) if cfg.get("runtime_probe_enabled", True) and isinstance(plan, dict): try: from runtime_probe import attach_runtime_probe, runtime_facts_line user_bits = [] for m in reversed(messages or []): if isinstance(m, dict) and m.get("role") == "user": c = m.get("content") user_bits.append(c if isinstance(c, str) else str(c or "")) break plan = attach_runtime_probe( plan, user_text=user_bits[0] if user_bits else "", cfg=cfg ) line = runtime_facts_line(plan) if line: log.info("%s", line) except Exception as exc: # noqa: BLE001 log.warning("runtime_probe skipped (stream): %s", exc) path_mode = bool(cfg.get("path_resolve_enabled", True)) and path_resolve_needed( plan ) if path_mode: executor = pick_path_resolve_model(cfg) else: executor = pick_agent_executor_model(cfg) timeout = float( cfg.get("path_resolve_timeout_sec") if path_mode else ( cfg.get("executor_timeout_sec") or cfg.get("call_timeout_sec") or 90 ) ) max_ctx = int(cfg.get("executor_input_chars", 12000)) use_stream = bool(cfg.get("executor_use_stream", False)) verbose = bool(cfg.get("progress_verbose", False)) from progress_ui import execution_banner yield sse_chunk( cid=cid, model="hierarchical", delta={ "content": execution_banner( path_mode=path_mode, plan=plan, model=executor ) }, ) if verbose: yield sse_chunk( cid=cid, model="hierarchical", delta={ "content": f"· модель `{executor}` · таймаут {int(timeout)}с\n" }, ) _ = inject_plan_context forward = prepare_agent_executor_forward( body, messages, plan, cfg, executor=executor, stream=use_stream, minimal=True, path_resolve=path_mode, ) n_tools = len(forward.get("tools") or []) approx = len(json.dumps(forward, ensure_ascii=False, default=str)) from agent_hier import context_fill_for_forward ctx_line, ctx_meta = context_fill_for_forward( forward, cfg, budget_chars=max_ctx ) log.info( "executor forward model=%s tools=%s bytes≈%s stream=%s %s", executor, n_tools, approx, use_stream, ctx_line, ) if verbose or cfg.get("show_context_fill", False): payload_line = f"- tools={n_tools}\n" if cfg.get("show_context_fill", False): payload_line += f"- {ctx_line}\n" yield sse_chunk( cid=cid, model="hierarchical", delta={"content": payload_line}, ) _ = ctx_meta try: body_timeout = float(forward.get("timeout") or timeout) client_timeout = body_timeout + 30.0 if use_stream: async with httpx.AsyncClient( timeout=httpx.Timeout(client_timeout, connect=15.0) ) as client: async with client.stream( "POST", f"{litellm_url.rstrip('/')}/v1/chat/completions", headers={ "Authorization": f"Bearer {litellm_key}", "Content-Type": "application/json", }, json=forward, ) as resp: if resp.status_code >= 400: err_body = (await resp.aread())[:300] raise RuntimeError( f"litellm HTTP {resp.status_code}: " f"{err_body.decode('utf-8', 'replace')}" ) buf = b"" aiter = resp.aiter_bytes().__aiter__() deadline = time.perf_counter() + client_timeout got_data = False while time.perf_counter() < deadline: try: piece = await asyncio.wait_for( aiter.__anext__(), timeout=12.0 ) except StopAsyncIteration: break except asyncio.TimeoutError: yield b": keepalive\n\n" if not got_data: yield sse_chunk( cid=cid, model="hierarchical", delta={ "content": "- executor: ждём LiteLLM…\n" }, ) continue got_data = True buf += piece while True: nl = buf.find(b"\n") if nl < 0: break line = buf[:nl].rstrip(b"\r") buf = buf[nl + 1 :] if not line: yield b"\n" continue yield line + b"\n" if buf.strip(): yield buf.rstrip(b"\r") + b"\n" yield b"data: [DONE]\n\n" else: # Non-stream: keepalive while awaiting Novita (headers otherwise block) async def _post() -> httpx.Response: async with httpx.AsyncClient( timeout=httpx.Timeout(client_timeout, connect=15.0) ) as client: return await client.post( f"{litellm_url.rstrip('/')}/v1/chat/completions", headers={ "Authorization": f"Bearer {litellm_key}", "Content-Type": "application/json", }, json=forward, ) task = asyncio.create_task(_post()) while not task.done(): try: await asyncio.wait_for(asyncio.shield(task), timeout=8.0) except asyncio.TimeoutError: yield b": keepalive\n\n" yield ( b": executor waiting LiteLLM non-stream\n\n" ) resp = task.result() if resp.status_code >= 400: raise RuntimeError( f"litellm HTTP {resp.status_code}: " f"{resp.text[:300]}" ) data = resp.json() msg = ((data.get("choices") or [{}])[0].get("message") or {}) has_tools = bool(msg.get("tool_calls")) if not has_tools: from agent_hier import ( completion_preview, force_kickstart_read, pick_agent_escalate_model, completion_has_tool_calls as _has_tc, ) log.warning( "executor no tool_calls after plan: %s", completion_preview(data), ) # 1) Escalate to DeepSeek with hard tool_choice if cfg.get("executor_no_tools_escalate", True): esc = pick_agent_escalate_model(cfg) esc_fwd = dict(forward) esc_fwd["model"] = esc esc_fwd["tool_choice"] = "required" path_hint = None try: from agent_hier import _first_path_hint path_hint = _first_path_hint( plan, forward.get("tools") or [], messages ) except Exception: # noqa: BLE001 path_hint = None nudge = ( "Call a tool now. Start with read_file" + (f" on `{path_hint}`." if path_hint else ".") + " No prose." ) esc_fwd["messages"] = list(esc_fwd.get("messages") or []) + [ {"role": "user", "content": nudge} ] yield ( f": no tool_calls → escalate `{esc}`\n\n".encode( "utf-8" ) ) async with httpx.AsyncClient( timeout=httpx.Timeout(client_timeout, connect=15.0) ) as client2: resp2 = await client2.post( f"{litellm_url.rstrip('/')}/v1/chat/completions", headers={ "Authorization": f"Bearer {litellm_key}", "Content-Type": "application/json", }, json=esc_fwd, ) if resp2.status_code < 400: data2 = resp2.json() if isinstance(data2, dict) and _has_tc(data2): data = data2 has_tools = True log.info("escalate no_tools → tool_calls OK") # 2) Kickstart read_file on plan path (keeps Zed loop alive) if not has_tools and cfg.get( "executor_synthetic_on_timeout", True ): yield b": no tool_calls -> kickstart read_file\n\n" if str(cfg.get("executor_synthetic_mode") or "").lower() in ( "never", ): # still kickstart once on first turn — otherwise dead end data = force_kickstart_read( plan=plan, tools=forward.get("tools") or [], model=executor, messages=messages, ) elif executor_fallback_completion: data = executor_fallback_completion( plan=plan, tools=forward.get("tools") or [], model=executor, messages=messages, cfg=cfg, reason="no_tool_calls", ) else: data = force_kickstart_read( plan=plan, tools=forward.get("tools") or [], model=executor, messages=messages, ) has_tools = bool( ( (data.get("choices") or [{}])[0].get("message") or {} ).get("tool_calls") ) # Visible status for Zed (short), details stay in SSE comments from progress_ui import plan_paths, tool_status tc_names: list[str] = [] msg0 = ((data.get("choices") or [{}])[0].get("message") or {}) for tc in msg0.get("tool_calls") or []: if isinstance(tc, dict): tc_names.append( str(((tc.get("function") or {}).get("name")) or "") ) paths = plan_paths(plan) yield sse_chunk( cid=cid, model="hierarchical", delta={ "content": tool_status( tc_names, path=paths[0] if paths else None ) }, ) yield ( f": executor OK tool_calls={has_tools}\n\n".encode("utf-8") ) for chunk in completion_to_sse_chunks( data, cid=cid, model=str(data.get("model") or executor) ): yield chunk except Exception as exc: # noqa: BLE001 err = (str(exc) or type(exc).__name__)[:220] log.exception("executor after plan failed: %s", err) if cfg.get("executor_synthetic_on_timeout", True): yield ( f": timeout → executor fallback ({err[:80]})\n\n".encode( "utf-8" ) ) if executor_fallback_completion: data = executor_fallback_completion( plan=plan, tools=forward.get("tools") or [], model=executor, messages=messages, cfg=cfg, reason=f"exception:{err[:80]}", ) else: data = synthetic_first_tool_completion( plan=plan, tools=forward.get("tools") or [], model=executor, messages=messages, max_synthetic=int( cfg.get("executor_max_synthetic_continues", 1) or 1 ), max_edit_failures=int( cfg.get("executor_max_edit_failures", 4) or 4 ), ) for chunk in completion_to_sse_chunks( data, cid=cid, model=executor ): yield chunk else: yield sse_chunk( cid=cid, model="hierarchical", delta={ "content": ( f"\nExecutor не ответил: `{err}`. " "Роутер жив — увеличь timeout / урежь контекст; " "повтор после ok обычно быстрее (cache).\n" ) }, ) yield sse_chunk( cid=cid, model="hierarchical", delta={}, finish_reason="stop" ) yield b"data: [DONE]\n\n" return # Confirm / cancel / amend — text plan for user content = hier.content or "" yield sse_chunk( cid=cid, model="hierarchical", delta={"content": f"\n---\n\n{content}"}, ) try: requests_metric.labels( tier=decision.tier.value, lane=decision.lane, model="hierarchical_agent", status="200", ).inc() except Exception: # noqa: BLE001 pass orchestrator.after_request( session_id, prompt_hash=prompt_hash, success=True, escalate=False ) yield sse_chunk( cid=cid, model="hierarchical", delta={}, finish_reason="stop" ) yield b"data: [DONE]\n\n" return StreamingResponse( gen(), media_type="text/event-stream", headers={ "X-Router-Meta": meta_header({**meta, "mode": "hierarchical_agent"}), "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }, )