aleksey
a2d238d92e
Make Zed Agent closer to Cursor: deterministic DevOps path index, live Traefik port probe before blind edits, stop-after-edit, and quieter Russian progress.
588 lines
24 KiB
Python
588 lines
24 KiB
Python
"""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",
|
||
},
|
||
)
|