Jelajahi Sumber

fix: validate chat message roles

zhenyu.hu 3 minggu lalu
induk
melakukan
d34add6e95

+ 11 - 5
docs/plans/todo-14-message-role-validation.md

@@ -1,6 +1,6 @@
 # Todo 14 Message Role Validation Plan
 
-**Status:** in_progress
+**Status:** done
 
 ## Goal
 
@@ -50,8 +50,14 @@ uv run pytest
 
 ## Verification
 
-Expected result:
+Result:
 
-- Request validation fails before provider/network work for invalid roles.
-- Runtime-generated tool replies still validate.
-- Full suite passes.
+- Red run passed as expected: invalid role and tool pre-message tests failed before implementation.
+- `uv run pytest tests/test_websocket_api.py tests/test_event_agent.py` passed: 21 tests, 1 existing Starlette deprecation warning.
+- `uv run pytest` passed: 34 tests, 1 existing Starlette deprecation warning.
+
+## Evaluation
+
+- `ChatMessage.role` now accepts only `system`, `user`, `assistant`, or `tool`.
+- `DebugRunRequest` rejects `tool` entries in `pre_messages`, so API callers cannot send orphan tool replies as front-loaded context.
+- Internally generated EventAgent tool replies still validate and serialize correctly.

+ 1 - 1
docs/plans/todos.md

@@ -40,6 +40,6 @@
 | 11 | done | `docs/plans/todo-11-prompt-workspace.md` | Improve prompt/pre-message editing with browser-side persistence and reusable prompt sets. | UI tests or focused JS tests prove prompts persist and can be restored. |
 | 12 | done | `docs/plans/todo-12-round-stats.md` | Move round statistics to structured backend events for token counts, cached tokens, TTFT, elapsed time, and per-turn summaries. | `uv run pytest tests/test_debug_runtime.py tests/test_websocket_api.py`; `uv run pytest`. |
 | 13 | done | `docs/plans/todo-13-real-queue-websocket-loop.md` | Make runtime queues real producer/consumer boundaries and let WebSocket run separate upstream/downstream tasks. | `uv run pytest tests/test_debug_runtime.py tests/test_websocket_api.py`; `uv run pytest`. |
-| 14 | in_progress | `docs/plans/todo-14-message-role-validation.md` | Reject provider-invalid API message roles and malformed tool replies before network calls. | Tests prove request validation rejects invalid roles and orphan tool messages. |
+| 14 | done | `docs/plans/todo-14-message-role-validation.md` | Reject provider-invalid API message roles and malformed tool replies before network calls. | `uv run pytest tests/test_websocket_api.py tests/test_event_agent.py`; `uv run pytest`. |
 | 15 | pending | `docs/plans/todo-15-static-package-data.md` | Include static UI assets in installed package builds. | Build/package inspection proves `index.html`, JS, and CSS are present. |
 | 16 | pending | `docs/plans/todo-16-tool-error-isolation.md` | Convert EventAgent tool handler exceptions into structured tool-result errors instead of aborting the session. | Tests prove a failing tool returns an error payload and the WebSocket run stays structured. |

+ 7 - 1
src/agent_lab/application/contracts.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from agent_lab.domain.messages import ChatMessage
 
@@ -26,3 +26,9 @@ class DebugRunRequest(BaseModel):
     pre_messages: list[ChatMessage] = Field(default_factory=list)
     chat_agent: AgentParams
     event_agent: EventAgentParams = Field(default_factory=EventAgentParams)
+
+    @model_validator(mode="after")
+    def reject_tool_pre_messages(self) -> "DebugRunRequest":
+        if any(message.role == "tool" for message in self.pre_messages):
+            raise ValueError("pre_messages cannot include tool messages")
+        return self

+ 2 - 2
src/agent_lab/domain/messages.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict
 
@@ -9,7 +9,7 @@ from agent_lab.domain.events import ToolCallEvent
 class ChatMessage(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    role: str
+    role: Literal["system", "user", "assistant", "tool"]
     content: str
     name: str | None = None
     tool_call_id: str | None = None

+ 37 - 0
tests/test_websocket_api.py

@@ -6,6 +6,7 @@ from pathlib import Path
 import httpx
 import pytest
 from fastapi.testclient import TestClient
+from pydantic import ValidationError
 
 from agent_lab.application.contracts import AgentParams, DebugRunRequest, EventAgentParams
 from agent_lab.application.queues import RuntimeQueues
@@ -153,6 +154,42 @@ def test_websocket_debug_sends_error_for_invalid_request():
     assert "user_message" in message["message"]
 
 
+def test_debug_run_request_rejects_invalid_pre_message_role():
+    payload = _request_payload()
+    payload["pre_messages"] = [{"role": "developer", "content": "invalid"}]
+
+    with pytest.raises(ValidationError) as exc_info:
+        DebugRunRequest.model_validate(payload)
+
+    assert "role" in str(exc_info.value)
+
+
+def test_debug_run_request_rejects_tool_pre_messages():
+    payload = _request_payload()
+    payload["pre_messages"] = [
+        {
+            "role": "tool",
+            "content": "orphan result",
+            "tool_call_id": "call_1",
+        }
+    ]
+
+    with pytest.raises(ValidationError, match="pre_messages cannot include tool messages"):
+        DebugRunRequest.model_validate(payload)
+
+
+def test_chat_message_allows_internal_tool_replies():
+    message = ChatMessage(
+        role="tool",
+        content='{"message":"handled"}',
+        name="handoff_note",
+        tool_call_id="call_1",
+    )
+
+    assert message.role == "tool"
+    assert message.tool_call_id == "call_1"
+
+
 class HistoryCapturingChatClient:
     def __init__(self) -> None:
         self.calls = 0