feat(agent): hierarchical executor with path resolve, runtime probe, quiet UI
Make Zed Agent closer to Cursor: deterministic DevOps path index, live Traefik port probe before blind edits, stop-after-edit, and quieter Russian progress.
This commit is contained in:
@@ -0,0 +1,587 @@
|
||||
"""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",
|
||||
},
|
||||
)
|
||||
Reference in New Issue
Block a user