Kaynağa Gözat

feat: reduce perceived reply latency

zhenyu.hu 3 hafta önce
ebeveyn
işleme
76e7f64dc3

+ 41 - 0
docs/plans/todo-19-latency-oriented-polish.md

@@ -0,0 +1,41 @@
+# Todo 19 Latency-Oriented Polish Plan
+
+**Status:** done
+
+## Goal
+
+Reduce perceived reply wait time with earlier event visibility and clearer frontend waiting feedback.
+
+## Design
+
+- Emit `event` output as soon as ChatAgent streaming produces an event, instead of waiting for the full model stream to finish.
+- Keep EventAgent execution after the assistant turn is finalized so history remains provider-compatible.
+- Show the user's submitted message immediately, disable the run button while streaming, and update elapsed time while waiting for the first token.
+- Add lightweight WebSocket logger entries for backend session boundaries.
+
+## Files
+
+- Modify: `src/agent_lab/application/runtime.py`
+- Modify: `src/agent_lab/presentation/web.py`
+- Modify: `src/agent_lab/presentation/static/index.html`
+- Modify: `src/agent_lab/presentation/static/app.js`
+- Modify: `src/agent_lab/presentation/static/styles.css`
+- Modify: `tests/test_debug_runtime.py`
+- Modify: `tests/test_websocket_api.py`
+- Modify: `docs/plans/todos.md`
+
+## Verification
+
+- Runtime test proves event messages reach the output queue before the ChatAgent stream is released.
+- WebSocket test proves backend logger records session lifecycle events.
+- Static UI tests prove immediate user echo, waiting state, disabled run button, and live elapsed timer hooks.
+- Full suite passes with `uv run pytest`.
+
+## Result
+
+- Runtime now emits `event` output immediately when the ChatAgent stream detects an event.
+- WebSocket service logs accepted, validated, terminal, rejected, failed, disconnected, and closed lifecycle states.
+- Frontend now echoes the submitted user message immediately, disables Run while active, shows "Waiting for first token", and keeps elapsed time moving before the first token.
+- Review follow-up: elapsed no longer freezes when an event-only first round emits `round_stats` before the first assistant token.
+- Tools management is moved to the top of the sidebar so it is visible in the first viewport.
+- Verified with focused tests, `uv run pytest` (`49 passed`), and local browser checks at `http://127.0.0.1:8001`.

+ 1 - 1
docs/plans/todos.md

@@ -45,4 +45,4 @@
 | 16 | done | `docs/plans/todo-16-tool-error-isolation.md` | Convert EventAgent tool handler exceptions into structured tool-result errors instead of aborting the session. | `uv run pytest tests/test_event_agent.py tests/test_debug_runtime.py`; `uv run pytest`. |
 | 17 | done | `docs/plans/todo-17-agent-config-audit-tools.md` | Align ChatAgent/EventAgent config, add extra_body defaults, default EventAgent to one event round, add mock tools, audit events, and backend logging. | `uv run pytest` passes (`42 passed`, one existing Starlette deprecation warning). |
 | 18 | done | `docs/plans/todo-18-prompt-modal-tool-ui.md` | Move prompt workspace into a modal and improve visible tool-management UI. | `uv run pytest` passes (`46 passed`), plus local browser check confirms modal and tool UI render. |
-| 19 | pending | `docs/plans/todo-19-latency-oriented-polish.md` | Do one optimization pass focused on reducing perceived reply wait time and cleaning frontend/backend rough edges. | Full test suite plus focused UI/static checks pass. |
+| 19 | done | `docs/plans/todo-19-latency-oriented-polish.md` | Do one optimization pass focused on reducing perceived reply wait time and cleaning frontend/backend rough edges. | `uv run pytest` passes (`49 passed`), plus local browser check confirms tools are visible in the first viewport. |

+ 3 - 4
src/agent_lab/application/runtime.py

@@ -155,6 +155,9 @@ class DebugRuntime:
                     saw_event = True
                     event = self._event_name_only(item.event)
                     events.append(event)
+                    await queues.output.put(
+                        {"type": "event", "event": event.model_dump()}
+                    )
                     await self._audit(
                         queues,
                         "chat_event_detected",
@@ -182,10 +185,6 @@ class DebugRuntime:
                 messages.append(assistant_message)
 
             if events:
-                for event in events:
-                    await queues.output.put(
-                        {"type": "event", "event": event.model_dump()}
-                    )
                 await queues.events.put(
                     EventAgentRequest(
                         events=events,

+ 41 - 7
src/agent_lab/presentation/static/app.js

@@ -2,6 +2,7 @@ const messagesEl = document.querySelector("#messages");
 const statusEl = document.querySelector("#connection-status");
 const chatForm = document.querySelector("#chat-form");
 const userMessage = document.querySelector("#user-message");
+const runButton = document.querySelector("#run-button");
 const promptSetName = document.querySelector("#prompt-set-name");
 const savedPromptSets = document.querySelector("#saved-prompt-sets");
 const systemPrompts = document.querySelector("#system-prompts");
@@ -16,6 +17,7 @@ let socket = null;
 let activeAssistant = null;
 let runStartedAt = 0;
 let firstTokenAt = 0;
+let elapsedTimer = 0;
 let hasBackendRoundStats = false;
 let pendingEnabledTools = null;
 
@@ -349,20 +351,32 @@ function writePromptSets(store) {
 }
 
 function runDebugSession() {
+  const submittedMessage = userMessage.value.trim();
+  if (!submittedMessage) {
+    statusEl.textContent = "Message required";
+    return;
+  }
+
   closeSocket();
   resetRun();
+  appendLog("user", submittedMessage);
+  setRunState(true);
+  runStartedAt = performance.now();
+  startElapsedTimer();
+  statusEl.textContent = "Connecting";
 
   socket = new WebSocket(wsUrl());
   socket.addEventListener("open", () => {
-    statusEl.textContent = "Streaming";
-    runStartedAt = performance.now();
-    socket.send(JSON.stringify(buildRequest()));
+    statusEl.textContent = "Waiting for first token";
+    socket.send(JSON.stringify(buildRequest(submittedMessage)));
   });
   socket.addEventListener("message", (event) => {
     handleServerMessage(JSON.parse(event.data));
   });
   socket.addEventListener("close", () => {
     statusEl.textContent = "Idle";
+    stopElapsedTimer();
+    setRunState(false);
     updateElapsed();
   });
   socket.addEventListener("error", () => {
@@ -370,9 +384,9 @@ function runDebugSession() {
   });
 }
 
-function buildRequest() {
+function buildRequest(submittedMessage = userMessage.value) {
   return {
-    user_message: userMessage.value,
+    user_message: submittedMessage,
     system_prompts: [...document.querySelectorAll(".system-prompt")]
       .map((input) => input.value.trim())
       .filter(Boolean),
@@ -456,6 +470,7 @@ function formatAuditMessage(message) {
 function appendAssistantDelta(content) {
   if (!firstTokenAt) {
     firstTokenAt = performance.now();
+    statusEl.textContent = "Streaming";
     if (!hasBackendRoundStats) {
       document.querySelector("#stat-ttft").textContent = `${Math.round(firstTokenAt - runStartedAt)}ms`;
     }
@@ -487,7 +502,9 @@ function updateUsage(usage) {
 }
 
 function updateRoundStats(stats) {
-  hasBackendRoundStats = true;
+  if (firstTokenAt) {
+    hasBackendRoundStats = true;
+  }
   document.querySelector("#stat-tokens").textContent = stats.total_tokens || 0;
   document.querySelector("#stat-cached").textContent = stats.cached_tokens || 0;
   const ttftText = formatMs(stats.ttft_ms);
@@ -505,13 +522,30 @@ function formatMs(value) {
 }
 
 function updateElapsed() {
-  if (!runStartedAt || hasBackendRoundStats) {
+  if (!runStartedAt || (hasBackendRoundStats && firstTokenAt)) {
     return;
   }
   document.querySelector("#stat-elapsed").textContent = `${Math.round(performance.now() - runStartedAt)}ms`;
 }
 
+function startElapsedTimer() {
+  stopElapsedTimer();
+  elapsedTimer = window.setInterval(updateElapsed, 100);
+}
+
+function stopElapsedTimer() {
+  if (elapsedTimer) {
+    window.clearInterval(elapsedTimer);
+    elapsedTimer = 0;
+  }
+}
+
+function setRunState(isRunning) {
+  runButton.disabled = isRunning;
+}
+
 function resetRun() {
+  stopElapsedTimer();
   activeAssistant = null;
   runStartedAt = 0;
   firstTokenAt = 0;

+ 12 - 11
src/agent_lab/presentation/static/index.html

@@ -17,6 +17,17 @@
 
     <main class="layout">
       <aside class="panel controls">
+        <section class="tools-section">
+          <div class="section-head">
+            <h2>Tools <span id="tool-count">0 selected</span></h2>
+            <div class="button-row compact">
+              <button id="select-all-tools" type="button">All</button>
+              <button id="clear-tools" type="button">None</button>
+            </div>
+          </div>
+          <div id="tool-list" class="stack">Loading tools...</div>
+        </section>
+
         <section>
           <h2>Chat Agent</h2>
           <label>Model <input id="model" placeholder="Backend default" /></label>
@@ -46,16 +57,6 @@
           </fieldset>
         </section>
 
-        <section>
-          <div class="section-head">
-            <h2>Tools <span id="tool-count">0 selected</span></h2>
-            <div class="button-row compact">
-              <button id="select-all-tools" type="button">All</button>
-              <button id="clear-tools" type="button">None</button>
-            </div>
-          </div>
-          <div id="tool-list" class="stack">Loading tools...</div>
-        </section>
       </aside>
 
       <section class="workspace">
@@ -70,7 +71,7 @@
 
         <form id="chat-form" class="composer">
           <textarea id="user-message" rows="3" placeholder="Send a debug message"></textarea>
-          <button type="submit">Run</button>
+          <button id="run-button" type="submit">Run</button>
         </form>
       </section>
     </main>

+ 11 - 0
src/agent_lab/presentation/static/styles.css

@@ -33,6 +33,11 @@ button:hover {
   background: #115e59;
 }
 
+button:disabled {
+  cursor: not-allowed;
+  opacity: 0.55;
+}
+
 .topbar {
   align-items: center;
   background: #ffffff;
@@ -194,6 +199,12 @@ textarea {
   border: 1px solid #b7ddda;
 }
 
+.user {
+  align-self: flex-end;
+  background: #eff6ff;
+  border: 1px solid #bfdbfe;
+}
+
 .session {
   background: #f4f6f8;
   color: #52616f;

+ 10 - 0
src/agent_lab/presentation/web.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any
@@ -21,6 +22,7 @@ RuntimeFactory = Callable[[], Any]
 
 STATIC_DIR = Path(__file__).parent / "static"
 TERMINAL_MESSAGE_TYPES = {"done", "error"}
+logger = logging.getLogger(__name__)
 
 
 def create_app(
@@ -52,11 +54,13 @@ def create_app(
     @app.websocket("/ws/debug")
     async def debug(websocket: WebSocket) -> None:
         await websocket.accept()
+        logger.info("websocket session accepted")
         runtime = resolved_runtime_factory()
         should_close = True
         try:
             payload = await websocket.receive_json()
             request = DebugRunRequest.model_validate(payload)
+            logger.info("websocket request accepted")
             if hasattr(runtime, "start"):
                 queues = runtime.start(request)
                 await _run_queue_session(websocket, queues)
@@ -64,14 +68,18 @@ def create_app(
                 async for message in runtime.run(request):
                     await websocket.send_json(message)
         except WebSocketDisconnect:
+            logger.info("websocket session disconnected")
             should_close = False
             return
         except (ValidationError, ValueError) as exc:
+            logger.warning("websocket request rejected: %s", exc)
             await websocket.send_json({"type": "error", "message": str(exc)})
         except Exception as exc:
+            logger.exception("websocket session failed")
             await websocket.send_json({"type": "error", "message": str(exc)})
         finally:
             await _close_runtime(runtime)
+            logger.info("websocket session closed")
             if should_close:
                 await websocket.close()
 
@@ -97,6 +105,7 @@ async def _send_downstream(websocket: WebSocket, queues: RuntimeQueues) -> None:
         message = await queues.output.get()
         await websocket.send_json(message)
         if message.get("type") in TERMINAL_MESSAGE_TYPES:
+            logger.info("websocket downstream terminal type=%s", message.get("type"))
             return
 
 
@@ -106,6 +115,7 @@ async def _receive_upstream(websocket: WebSocket, queues: RuntimeQueues) -> None
         try:
             message = _parse_upstream_message(payload)
         except (TypeError, ValueError) as exc:
+            logger.warning("websocket upstream rejected: %s", exc)
             await queues.output.put({"type": "error", "message": str(exc)})
             return
         await queues.input.put(message)

+ 72 - 0
tests/test_debug_runtime.py

@@ -41,6 +41,13 @@ async def _next_non_audit(
             return message
 
 
+async def _next_non_audit_from_queue(queues: Any) -> dict[str, Any]:
+    while True:
+        message = await queues.output.get()
+        if message["type"] != "audit":
+            return message
+
+
 class RecordingQueue(asyncio.Queue):
     def __init__(self, name: str, log: list[tuple[str, str, str]]) -> None:
         super().__init__()
@@ -250,6 +257,36 @@ class EventRoundStatsChatClient:
         )
 
 
+class SlowAfterEventChatClient:
+    def __init__(self) -> None:
+        self.calls = 0
+        self.event_seen = asyncio.Event()
+        self.release_stream = asyncio.Event()
+
+    async def stream_chat(
+        self,
+        messages: list[ChatMessage],
+        tools: list[dict],
+        params: AgentParams,
+    ) -> AsyncIterator[StreamItem]:
+        self.calls += 1
+        if self.calls == 1:
+            yield StreamItem.message_delta("Need event.")
+            yield StreamItem.event(
+                ToolCallEvent(
+                    id="call_1",
+                    name="handoff_note",
+                    arguments={},
+                    raw_arguments="{}",
+                )
+            )
+            self.event_seen.set()
+            await self.release_stream.wait()
+            return
+
+        yield StreamItem.message_delta("final answer")
+
+
 def test_runtime_queues_exposes_input_output_and_events_queues():
     RuntimeQueues = _runtime_queues_class()
 
@@ -325,6 +362,41 @@ async def test_runtime_emits_audit_events_and_backend_logs(caplog):
     assert "event_agent_completed" in caplog.text
 
 
+@pytest.mark.asyncio
+async def test_runtime_outputs_event_as_soon_as_chat_stream_detects_it():
+    RuntimeQueues = _runtime_queues_class()
+    queues = RuntimeQueues()
+    client = SlowAfterEventChatClient()
+    runtime = DebugRuntime(client, queues=queues)
+    request = DebugRunRequest(
+        user_message="debug this",
+        system_prompts=[],
+        pre_messages=[],
+        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
+        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
+    )
+
+    runtime.start(request)
+
+    try:
+        assert await _next_non_audit_from_queue(queues) == {"type": "session_started"}
+        assert await _next_non_audit_from_queue(queues) == {
+            "type": "message_delta",
+            "content": "Need event.",
+        }
+        await asyncio.wait_for(client.event_seen.wait(), timeout=1)
+        event_message = await asyncio.wait_for(
+            _next_non_audit_from_queue(queues),
+            timeout=0.2,
+        )
+
+        assert event_message["type"] == "event"
+        assert event_message["event"]["name"] == "handoff_note"
+    finally:
+        client.release_stream.set()
+        await runtime.aclose()
+
+
 @pytest.mark.asyncio
 async def test_runtime_batches_round_events_before_continuing_chat_agent():
     def resolve_from_history(

+ 38 - 0
tests/test_websocket_api.py

@@ -1,5 +1,6 @@
 import json
 import asyncio
+import logging
 from collections.abc import AsyncIterator
 from pathlib import Path
 
@@ -131,6 +132,26 @@ def test_websocket_debug_streams_runtime_messages():
     assert runtime.requests[0].pre_messages[0].content == "previous turn"
 
 
+def test_websocket_debug_logs_session_lifecycle(caplog):
+    runtime = FakeRuntime()
+    app = create_app(runtime_factory=lambda: runtime)
+    client = TestClient(app)
+
+    with caplog.at_level(logging.INFO, logger="agent_lab.presentation.web"):
+        with client.websocket_connect("/ws/debug") as websocket:
+            websocket.send_json(_request_payload())
+            assert websocket.receive_json() == {"type": "session_started"}
+            assert websocket.receive_json() == {
+                "type": "message_delta",
+                "content": "hello",
+            }
+            assert websocket.receive_json() == {"type": "done"}
+
+    assert "websocket session accepted" in caplog.text
+    assert "websocket request accepted" in caplog.text
+    assert "websocket session closed" in caplog.text
+
+
 def test_websocket_debug_enqueues_user_messages_during_running_session():
     runtime = QueueAwareRuntime()
     app = create_app(runtime_factory=lambda: runtime)
@@ -526,6 +547,23 @@ def test_static_app_handles_audit_messages():
     assert ".audit" in css
 
 
+def test_static_app_shows_wait_state_and_immediate_user_echo():
+    html = Path("src/agent_lab/presentation/static/index.html").read_text()
+    js = Path("src/agent_lab/presentation/static/app.js").read_text()
+    css = Path("src/agent_lab/presentation/static/styles.css").read_text()
+
+    assert 'id="run-button"' in html
+    assert 'const runButton = document.querySelector("#run-button")' in js
+    assert 'appendLog("user", submittedMessage)' in js
+    assert 'statusEl.textContent = "Waiting for first token"' in js
+    assert "function startElapsedTimer()" in js
+    assert "function stopElapsedTimer()" in js
+    assert "function setRunState(" in js
+    assert "if (firstTokenAt) {\n    hasBackendRoundStats = true;" in js
+    assert "hasBackendRoundStats && firstTokenAt" in js
+    assert ".user" in css
+
+
 def test_static_prompt_workspace_uses_stable_storage_hooks():
     js = Path("src/agent_lab/presentation/static/app.js").read_text()