290 lines
13 KiB
Go
290 lines
13 KiB
Go
package handler
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"fire-safety-ymd/internal/service"
|
|
)
|
|
|
|
const testChatToken = "abcdef0123456789abcdef0123456789"
|
|
|
|
func TestNewChatHandlerRejectsUnsafeOptions(t *testing.T) {
|
|
tests := []ChatOptions{
|
|
{AuthToken: "short", MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
{AuthToken: "", AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
{AuthToken: "delivered\tkey", AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
{AuthToken: "delivered密钥", AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
{AuthToken: strings.Repeat("a", 4097), AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
{AuthToken: testChatToken, MaxBodyBytes: maximumChatBodyBytes + 1, RunTimeout: time.Second},
|
|
{AuthToken: testChatToken, MaxBodyBytes: 1, RunTimeout: maximumChatRunTime + time.Second},
|
|
{AuthToken: testChatToken, AllowedOrigins: []string{"*"}, MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
{AuthToken: testChatToken, AllowedOrigins: []string{"https://fire.example.test/"}, MaxBodyBytes: 1, RunTimeout: time.Second},
|
|
}
|
|
for _, options := range tests {
|
|
if _, err := NewChatHandler(&fakeChatUseCase{}, options); err == nil {
|
|
t.Fatalf("NewChatHandler(%#v) error = nil", options)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestNewChatHandlerAllowsExplicitLegacyShortToken(t *testing.T) {
|
|
_, err := NewChatHandler(&fakeChatUseCase{}, ChatOptions{
|
|
AuthToken: "delivered-key",
|
|
AllowLegacyShortToken: true,
|
|
MaxBodyBytes: 1,
|
|
RunTimeout: time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewChatHandler() error = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestChatTransportGuardsRunBeforeService(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
method string
|
|
body string
|
|
configure func(*http.Request)
|
|
maxBody int64
|
|
wantStatus int
|
|
wantCode string
|
|
}{
|
|
{name: "wrong method", method: http.MethodGet, body: `{}`, wantStatus: http.StatusMethodNotAllowed, wantCode: "CHAT_METHOD_NOT_ALLOWED"},
|
|
{name: "forbidden origin", method: http.MethodPost, body: `{"message":"hello"}`, configure: func(r *http.Request) { r.Header.Set("Origin", "https://evil.example") }, wantStatus: http.StatusForbidden, wantCode: "CHAT_ORIGIN_FORBIDDEN"},
|
|
{name: "missing auth", method: http.MethodPost, body: `{"message":"hello"}`, configure: func(r *http.Request) { r.Header.Del("Authorization") }, wantStatus: http.StatusUnauthorized, wantCode: "CHAT_AUTH_INVALID"},
|
|
{name: "wrong content type", method: http.MethodPost, body: `{"message":"hello"}`, configure: func(r *http.Request) { r.Header.Set("Content-Type", "text/plain") }, wantStatus: http.StatusUnsupportedMediaType, wantCode: "CHAT_CONTENT_TYPE_INVALID"},
|
|
{name: "wrong accept", method: http.MethodPost, body: `{"message":"hello"}`, configure: func(r *http.Request) { r.Header.Set("Accept", "application/json") }, wantStatus: http.StatusNotAcceptable, wantCode: "CHAT_ACCEPT_INVALID"},
|
|
{name: "event stream explicitly rejected", method: http.MethodPost, body: `{"message":"hello"}`, configure: func(r *http.Request) { r.Header.Set("Accept", "text/event-stream;q=0") }, wantStatus: http.StatusNotAcceptable, wantCode: "CHAT_ACCEPT_INVALID"},
|
|
{name: "invalid JSON", method: http.MethodPost, body: `{"message":`, wantStatus: http.StatusBadRequest, wantCode: "CHAT_REQUEST_INVALID"},
|
|
{name: "unknown field", method: http.MethodPost, body: `{"message":"hello","user_id":"admin"}`, wantStatus: http.StatusBadRequest, wantCode: "CHAT_REQUEST_INVALID"},
|
|
{name: "missing message", method: http.MethodPost, body: `{}`, wantStatus: http.StatusBadRequest, wantCode: "CHAT_REQUEST_INVALID"},
|
|
{name: "null conversation", method: http.MethodPost, body: `{"message":"hello","conversation_id":null}`, wantStatus: http.StatusBadRequest, wantCode: "CHAT_REQUEST_INVALID"},
|
|
{name: "empty conversation", method: http.MethodPost, body: `{"message":"hello","conversation_id":""}`, wantStatus: http.StatusBadRequest, wantCode: "CHAT_REQUEST_INVALID"},
|
|
{name: "body too large", method: http.MethodPost, body: `{"message":"hello"}`, maxBody: 4, wantStatus: http.StatusRequestEntityTooLarge, wantCode: "CHAT_REQUEST_BODY_TOO_LARGE"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
chat := &fakeChatUseCase{}
|
|
maxBody := tt.maxBody
|
|
if maxBody == 0 {
|
|
maxBody = 4096
|
|
}
|
|
handler := newTestChatHandler(t, chat, maxBody, []string{"https://allowed.example"}, nil)
|
|
request := authenticatedChatRequest(tt.method, tt.body)
|
|
if tt.configure != nil {
|
|
tt.configure(request)
|
|
}
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != tt.wantStatus || !strings.Contains(response.Body.String(), `"code":"`+tt.wantCode+`"`) {
|
|
t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
|
|
}
|
|
if chat.prepareCalls != 0 {
|
|
t.Fatalf("Prepare calls = %d, want 0", chat.prepareCalls)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestChatCORSPreflight(t *testing.T) {
|
|
handler := newTestChatHandler(t, &fakeChatUseCase{}, 4096, []string{"http://localhost:5173"}, nil)
|
|
|
|
allowed := httptest.NewRequest(http.MethodOptions, "/api/chat", nil)
|
|
allowed.Header.Set("Origin", "http://localhost:5173")
|
|
allowed.Header.Set("Access-Control-Request-Method", http.MethodPost)
|
|
allowed.Header.Set("Access-Control-Request-Headers", "authorization, content-type")
|
|
allowedResponse := httptest.NewRecorder()
|
|
handler.ServeHTTP(allowedResponse, allowed)
|
|
if allowedResponse.Code != http.StatusNoContent || allowedResponse.Header().Get("Access-Control-Allow-Origin") != "http://localhost:5173" {
|
|
t.Fatalf("allowed preflight status=%d headers=%#v", allowedResponse.Code, allowedResponse.Header())
|
|
}
|
|
|
|
forbidden := httptest.NewRequest(http.MethodOptions, "/api/chat", nil)
|
|
forbidden.Header.Set("Origin", "http://localhost:5173")
|
|
forbidden.Header.Set("Access-Control-Request-Method", http.MethodPost)
|
|
forbidden.Header.Set("Access-Control-Request-Headers", "x-admin-role")
|
|
forbiddenResponse := httptest.NewRecorder()
|
|
handler.ServeHTTP(forbiddenResponse, forbidden)
|
|
if forbiddenResponse.Code != http.StatusForbidden || !strings.Contains(forbiddenResponse.Body.String(), "CHAT_ORIGIN_FORBIDDEN") {
|
|
t.Fatalf("forbidden preflight status=%d body=%s", forbiddenResponse.Code, forbiddenResponse.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestChatStreamsSafeProgressAndFinalAnswer(t *testing.T) {
|
|
turn := &fakeChatTurn{
|
|
conversationID: "conv_0123456789abcdef01234567",
|
|
events: []service.AgentTraceEvent{
|
|
{Event: "message.delta", Status: "running"},
|
|
{Event: "tool.started", ToolName: "fire_safety_find_nearby_water_sources", Status: "running"},
|
|
},
|
|
result: service.ChatResult{
|
|
ConversationID: "conv_0123456789abcdef01234567",
|
|
RunID: "run-safe-1",
|
|
Answer: "候选水源需要现场确认。",
|
|
Usage: service.ChatTokenUsage{Input: 10, Output: 20, Total: 30},
|
|
},
|
|
}
|
|
chat := &fakeChatUseCase{turn: turn}
|
|
var logs bytes.Buffer
|
|
handler := newTestChatHandler(t, chat, 4096, []string{"https://allowed.example"}, log.New(&logs, "", 0))
|
|
request := authenticatedChatRequest(http.MethodPost, `{"message":"不要记录这个问题 secret-user-text"}`)
|
|
request.Header.Set("Origin", "https://allowed.example")
|
|
request.Header.Set("X-Request-ID", "chat-request-1")
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
if response.Code != http.StatusOK || !strings.HasPrefix(response.Header().Get("Content-Type"), "text/event-stream") {
|
|
t.Fatalf("status=%d content-type=%q body=%s", response.Code, response.Header().Get("Content-Type"), response.Body.String())
|
|
}
|
|
body := response.Body.String()
|
|
for _, expected := range []string{"event: conversation", "event: progress", "event: message", "event: done", "候选水源需要现场确认。", `"reused":false`} {
|
|
if !strings.Contains(body, expected) {
|
|
t.Fatalf("SSE body missing %q: %s", expected, body)
|
|
}
|
|
}
|
|
if strings.Contains(body, "message.delta") || strings.Contains(body, "provider-session") {
|
|
t.Fatalf("SSE body exposed suppressed provider details: %s", body)
|
|
}
|
|
if !(strings.Index(body, "event: conversation") < strings.Index(body, "event: progress") &&
|
|
strings.Index(body, "event: progress") < strings.Index(body, "event: message") &&
|
|
strings.Index(body, "event: message") < strings.Index(body, "event: done")) {
|
|
t.Fatalf("unexpected SSE order: %s", body)
|
|
}
|
|
if chat.lastRequest.Message != "不要记录这个问题 secret-user-text" || chat.lastRequest.RequestID != "chat-request-1" {
|
|
t.Fatalf("service request = %#v", chat.lastRequest)
|
|
}
|
|
if !turn.closed {
|
|
t.Fatal("turn was not closed")
|
|
}
|
|
if strings.Contains(logs.String(), "secret-user-text") || strings.Contains(logs.String(), "候选水源") {
|
|
t.Fatalf("logs contain chat content: %s", logs.String())
|
|
}
|
|
}
|
|
|
|
func TestChatStreamFailureDoesNotExposePartialAnswer(t *testing.T) {
|
|
turn := &fakeChatTurn{
|
|
conversationID: "conv_0123456789abcdef01234567",
|
|
events: []service.AgentTraceEvent{{Event: "tool.failed", ToolName: "safe_tool", Status: "failed"}},
|
|
err: fmt.Errorf("provider secret response: %w", service.ErrChatUpstreamProtocol),
|
|
}
|
|
handler := newTestChatHandler(t, &fakeChatUseCase{turn: turn}, 4096, nil, nil)
|
|
request := authenticatedChatRequest(http.MethodPost, `{"message":"hello"}`)
|
|
response := httptest.NewRecorder()
|
|
|
|
handler.ServeHTTP(response, request)
|
|
|
|
body := response.Body.String()
|
|
if response.Code != http.StatusOK || !strings.Contains(body, "event: error") || !strings.Contains(body, "CHAT_UPSTREAM_PROTOCOL_ERROR") {
|
|
t.Fatalf("status=%d body=%s", response.Code, body)
|
|
}
|
|
if strings.Contains(body, "event: message") || strings.Contains(body, "event: done") || strings.Contains(body, "provider secret response") {
|
|
t.Fatalf("failure stream exposed final/secret data: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestChatMapsPreparationErrorsBeforeStartingSSE(t *testing.T) {
|
|
tests := []struct {
|
|
err error
|
|
wantStatus int
|
|
wantCode string
|
|
}{
|
|
{err: service.ErrChatInvalidArgument, wantStatus: http.StatusBadRequest, wantCode: "CHAT_REQUEST_INVALID"},
|
|
{err: service.ErrChatConversationNotFound, wantStatus: http.StatusNotFound, wantCode: "CHAT_CONVERSATION_NOT_FOUND"},
|
|
{err: service.ErrChatConversationBusy, wantStatus: http.StatusConflict, wantCode: "CHAT_CONVERSATION_BUSY"},
|
|
{err: service.ErrChatCapacityReached, wantStatus: http.StatusServiceUnavailable, wantCode: "CHAT_CAPACITY_REACHED"},
|
|
{err: context.DeadlineExceeded, wantStatus: http.StatusGatewayTimeout, wantCode: "CHAT_UPSTREAM_TIMEOUT"},
|
|
{err: service.ErrChatUpstreamUnavailable, wantStatus: http.StatusBadGateway, wantCode: "CHAT_UPSTREAM_UNAVAILABLE"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.wantCode, func(t *testing.T) {
|
|
handler := newTestChatHandler(t, &fakeChatUseCase{err: tt.err}, 4096, nil, nil)
|
|
request := authenticatedChatRequest(http.MethodPost, `{"message":"hello"}`)
|
|
response := httptest.NewRecorder()
|
|
handler.ServeHTTP(response, request)
|
|
if response.Code != tt.wantStatus || !strings.Contains(response.Body.String(), tt.wantCode) || strings.Contains(response.Header().Get("Content-Type"), "text/event-stream") {
|
|
t.Fatalf("status=%d content-type=%q body=%s", response.Code, response.Header().Get("Content-Type"), response.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func newTestChatHandler(t *testing.T, chat ChatUseCase, maxBody int64, origins []string, logger *log.Logger) *ChatHandler {
|
|
t.Helper()
|
|
handler, err := NewChatHandler(chat, ChatOptions{
|
|
AuthToken: testChatToken,
|
|
AllowedOrigins: origins,
|
|
MaxBodyBytes: maxBody,
|
|
RunTimeout: time.Minute,
|
|
Logger: logger,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewChatHandler() error = %v", err)
|
|
}
|
|
return handler
|
|
}
|
|
|
|
func authenticatedChatRequest(method, body string) *http.Request {
|
|
request := httptest.NewRequest(method, "/api/chat", strings.NewReader(body))
|
|
request.Header.Set("Authorization", "Bearer "+testChatToken)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
request.Header.Set("Accept", "text/event-stream")
|
|
return request
|
|
}
|
|
|
|
type fakeChatUseCase struct {
|
|
turn service.ChatTurn
|
|
err error
|
|
prepareCalls int
|
|
lastRequest service.ChatRequest
|
|
}
|
|
|
|
func (f *fakeChatUseCase) Prepare(_ context.Context, request service.ChatRequest) (service.ChatTurn, error) {
|
|
f.prepareCalls++
|
|
f.lastRequest = request
|
|
if f.err != nil {
|
|
return nil, f.err
|
|
}
|
|
if f.turn == nil {
|
|
return &fakeChatTurn{conversationID: "conv_0123456789abcdef01234567", result: service.ChatResult{
|
|
ConversationID: "conv_0123456789abcdef01234567",
|
|
Answer: "ok",
|
|
}}, nil
|
|
}
|
|
return f.turn, nil
|
|
}
|
|
|
|
type fakeChatTurn struct {
|
|
conversationID string
|
|
reused bool
|
|
events []service.AgentTraceEvent
|
|
result service.ChatResult
|
|
err error
|
|
closed bool
|
|
}
|
|
|
|
func (f *fakeChatTurn) ConversationID() string { return f.conversationID }
|
|
|
|
func (f *fakeChatTurn) Reused() bool { return f.reused }
|
|
|
|
func (f *fakeChatTurn) Stream(_ context.Context, trace func(service.AgentTraceEvent)) (service.ChatResult, error) {
|
|
for _, event := range f.events {
|
|
if trace != nil {
|
|
trace(event)
|
|
}
|
|
}
|
|
return f.result, f.err
|
|
}
|
|
|
|
func (f *fakeChatTurn) Close() { f.closed = true }
|