fire-safety-ymd/internal/handler/dashscope_chat_test.go

package handler

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"fire-safety-ymd/internal/service"
)

const testDashScopeAppID = "fire-safety-test-app"

func TestNewDashScopeChatHandlerRejectsUnsafeOptions(t *testing.T) {
	tests := []DashScopeChatOptions{
		{AppID: "", AuthToken: testChatToken, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: "app/other", AuthToken: testChatToken, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: "app.other", AuthToken: testChatToken, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: "short", MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: "", AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: "delivered\tkey", AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: "delivered密钥", AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: strings.Repeat("a", 4097), AllowLegacyShortToken: true, MaxBodyBytes: 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: testChatToken, MaxBodyBytes: maximumChatBodyBytes + 1, RunTimeout: time.Second},
		{AppID: testDashScopeAppID, AuthToken: testChatToken, MaxBodyBytes: 1, RunTimeout: maximumChatRunTime + time.Second},
		{AppID: testDashScopeAppID, AuthToken: testChatToken, AllowedOrigins: []string{"*"}, MaxBodyBytes: 1, RunTimeout: time.Second},
	}
	for _, options := range tests {
		if _, err := NewDashScopeChatHandler(&fakeChatUseCase{}, options); err == nil {
			t.Fatalf("NewDashScopeChatHandler(%#v) error = nil", options)
		}
	}
	if _, err := NewDashScopeChatHandler(nil, DashScopeChatOptions{
		AppID: testDashScopeAppID, AuthToken: testChatToken, MaxBodyBytes: 1, RunTimeout: time.Second,
	}); err == nil {
		t.Fatal("NewDashScopeChatHandler(nil) error = nil")
	}
}

func TestNewDashScopeChatHandlerAllowsExplicitLegacyShortToken(t *testing.T) {
	_, err := NewDashScopeChatHandler(&fakeChatUseCase{}, DashScopeChatOptions{
		AppID:                 testDashScopeAppID,
		AuthToken:             "delivered-key",
		AllowLegacyShortToken: true,
		MaxBodyBytes:          1,
		RunTimeout:            time.Second,
	})
	if err != nil {
		t.Fatalf("NewDashScopeChatHandler() error = %v", err)
	}
}

func TestDashScopeChatStreamsCompatibleResultAfterStrictSuccess(t *testing.T) {
	turn := &fakeChatTurn{
		conversationID: "conv_0123456789abcdef01234567",
		events: []service.AgentTraceEvent{
			{Event: "message.delta", Status: "running"},
			{Event: "tool.started", ToolName: "sensitive_tool", Status: "running"},
		},
		result: service.ChatResult{
			ConversationID: "conv_0123456789abcdef01234567",
			RunID:          "provider-run-must-not-leak",
			ModelID:        "qwen-plus-latest",
			Answer:         "候选水源需要现场确认。",
			Usage:          service.ChatTokenUsage{Input: 12, Output: 34, Total: 46},
		},
	}
	chat := &fakeChatUseCase{turn: turn}
	var logs bytes.Buffer
	handler := newTestDashScopeChatHandler(t, chat, []string{"https://allowed.example"}, log.New(&logs, "", 0))
	request := authenticatedDashScopeRequest(http.MethodPost, `{"input":{"prompt":"不要记录 secret prompt"},"parameters":{"incremental_output":true},"debug":{}}`)
	request.Header.Set("Origin", "https://allowed.example")
	request.Header.Set("X-Request-ID", "compat-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())
	}
	if response.Header().Get("X-Accel-Buffering") != "no" || response.Header().Get("Cache-Control") != "no-store" {
		t.Fatalf("stream headers = %#v", response.Header())
	}
	blocks := dashScopeSSEBlocks(response.Body.String())
	if len(blocks) != 2 {
		t.Fatalf("SSE block count=%d body=%s", len(blocks), response.Body.String())
	}
	for index, expectedID := range []string{"id: 1", "id: 2"} {
		if !strings.Contains(blocks[index], expectedID) || !strings.Contains(blocks[index], "event: result") || !strings.Contains(blocks[index], ":HTTP_STATUS/200") {
			t.Fatalf("block %d is incompatible: %s", index, blocks[index])
		}
	}
	initial := decodeDashScopeSSEData(t, blocks[0])
	if initial.Output.SessionID != turn.conversationID || initial.Output.FinishReason != "null" || initial.Output.Text != "" || len(initial.Usage.Models) != 0 || initial.RequestID != "compat-request-1" {
		t.Fatalf("initial payload = %#v", initial)
	}
	final := decodeDashScopeSSEData(t, blocks[1])
	if final.Output.SessionID != turn.conversationID || final.Output.FinishReason != "stop" || final.Output.Text != "候选水源需要现场确认。" || final.RequestID != "compat-request-1" {
		t.Fatalf("final payload = %#v", final)
	}
	if len(final.Usage.Models) != 1 || final.Usage.Models[0].InputTokens != 12 || final.Usage.Models[0].OutputTokens != 34 || final.Usage.Models[0].ModelID != "qwen-plus-latest" {
		t.Fatalf("final usage = %#v", final.Usage)
	}
	body := response.Body.String()
	for _, forbidden := range []string{"message.delta", "sensitive_tool", "provider-run-must-not-leak"} {
		if strings.Contains(body, forbidden) {
			t.Fatalf("SSE exposed %q: %s", forbidden, body)
		}
	}
	if chat.lastRequest.Message != "不要记录 secret prompt" || chat.lastRequest.ConversationID != "" || chat.lastRequest.RequestID != "compat-request-1" {
		t.Fatalf("service request = %#v", chat.lastRequest)
	}
	if !turn.closed {
		t.Fatal("turn was not closed")
	}
	for _, sensitive := range []string{"secret prompt", "候选水源", turn.conversationID, testChatToken} {
		if strings.Contains(logs.String(), sensitive) {
			t.Fatalf("logs contain sensitive value %q: %s", sensitive, logs.String())
		}
	}
}

func TestDashScopeChatMapsSessionIDToLocalConversation(t *testing.T) {
	turn := &fakeChatTurn{
		conversationID: "conv_0123456789abcdef01234567",
		reused:         true,
		result: service.ChatResult{
			ConversationID: "conv_0123456789abcdef01234567",
			Answer:         "ok",
		},
	}
	chat := &fakeChatUseCase{turn: turn}
	handler := newTestDashScopeChatHandler(t, chat, nil, nil)
	request := authenticatedDashScopeRequest(http.MethodPost, `{"input":{"prompt":"follow up","session_id":"conv_0123456789abcdef01234567"},"parameters":{}}`)
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
	}
	if chat.lastRequest.ConversationID != "conv_0123456789abcdef01234567" {
		t.Fatalf("conversation ID = %q", chat.lastRequest.ConversationID)
	}
}

func TestDashScopeChatTransportGuardsBeforeService(t *testing.T) {
	tests := []struct {
		name       string
		method     string
		body       string
		configure  func(*http.Request)
		maxBody    int64
		wantStatus int
		wantCode   string
	}{
		{name: "wrong app", method: http.MethodPost, body: `{"input":{"prompt":"hello"}}`, configure: func(r *http.Request) { r.SetPathValue("app_id", "another-app") }, wantStatus: http.StatusNotFound, wantCode: "CHAT_APP_NOT_FOUND"},
		{name: "wrong method", method: http.MethodGet, body: `{}`, wantStatus: http.StatusMethodNotAllowed, wantCode: "CHAT_METHOD_NOT_ALLOWED"},
		{name: "forbidden origin", method: http.MethodPost, body: `{"input":{"prompt":"hello"}}`, configure: func(r *http.Request) { r.Header.Set("Origin", "https://evil.example") }, wantStatus: http.StatusForbidden, wantCode: "CHAT_ORIGIN_FORBIDDEN"},
		{name: "missing xtoken", method: http.MethodPost, body: `{"input":{"prompt":"hello"}}`, configure: func(r *http.Request) { r.Header.Del("xtoken") }, wantStatus: http.StatusUnauthorized, wantCode: "CHAT_AUTH_INVALID"},
		{name: "bearer is not xtoken", method: http.MethodPost, body: `{"input":{"prompt":"hello"}}`, configure: func(r *http.Request) { r.Header.Del("xtoken"); r.Header.Set("Authorization", "Bearer "+testChatToken) }, wantStatus: http.StatusUnauthorized, wantCode: "CHAT_AUTH_INVALID"},
		{name: "wrong content type", method: http.MethodPost, body: `{"input":{"prompt":"hello"}}`, configure: func(r *http.Request) { r.Header.Set("Content-Type", "text/plain") }, wantStatus: http.StatusUnsupportedMediaType, wantCode: "CHAT_CONTENT_TYPE_INVALID"},
		{name: "body too large", method: http.MethodPost, body: `{"input":{"prompt":"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 := newTestDashScopeChatHandlerWithLimit(t, chat, maxBody, []string{"https://allowed.example"}, nil)
			request := authenticatedDashScopeRequest(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 TestDashScopeChatCORSPreflight(t *testing.T) {
	handler := newTestDashScopeChatHandler(t, &fakeChatUseCase{}, []string{"http://localhost:5173"}, nil)

	allowed := authenticatedDashScopeRequest(http.MethodOptions, "")
	allowed.Header.Set("Origin", "http://localhost:5173")
	allowed.Header.Set("Access-Control-Request-Method", http.MethodPost)
	allowed.Header.Set("Access-Control-Request-Headers", "content-type, xtoken, x-dashscope-sse, x-request-id")
	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 body=%s", allowedResponse.Code, allowedResponse.Header(), allowedResponse.Body.String())
	}

	forbidden := authenticatedDashScopeRequest(http.MethodOptions, "")
	forbidden.Header.Set("Origin", "http://localhost:5173")
	forbidden.Header.Set("Access-Control-Request-Method", http.MethodPost)
	forbidden.Header.Set("Access-Control-Request-Headers", "authorization")
	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 TestDashScopeChatRejectsNonCompatibleJSON(t *testing.T) {
	tests := []string{
		``,
		`[]`,
		`{}`,
		`{"input":null}`,
		`{"input":{}}`,
		`{"input":{"prompt":null}}`,
		`{"input":{"prompt":" "}}`,
		`{"input":{"prompt":"hello","session_id":null}}`,
		`{"input":{"prompt":"hello","user_id":"admin"}}`,
		`{"input":{"prompt":"hello"},"parameters":null}`,
		`{"input":{"prompt":"hello"},"parameters":{"incremental_output":null}}`,
		`{"input":{"prompt":"hello"},"parameters":{"temperature":1}}`,
		`{"input":{"prompt":"hello"},"parameters":{"incremental_output":"yes"}}`,
		`{"input":{"prompt":"hello"},"debug":null}`,
		`{"input":{"prompt":"hello"},"debug":{"trace":true}}`,
		`{"input":{"prompt":"hello"},"metadata":{"role":"admin"}}`,
		`{"input":{"prompt":"hello"}} {}`,
	}
	for _, body := range tests {
		t.Run(fmt.Sprintf("%q", body), func(t *testing.T) {
			chat := &fakeChatUseCase{}
			handler := newTestDashScopeChatHandler(t, chat, nil, nil)
			request := authenticatedDashScopeRequest(http.MethodPost, body)
			response := httptest.NewRecorder()

			handler.ServeHTTP(response, request)

			if response.Code != http.StatusBadRequest || !strings.Contains(response.Body.String(), `"code":"CHAT_REQUEST_INVALID"`) {
				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 TestDashScopeChatStreamFailureDoesNotExposePartialAnswerOrStop(t *testing.T) {
	turn := &fakeChatTurn{
		conversationID: "conv_0123456789abcdef01234567",
		events: []service.AgentTraceEvent{
			{Event: "message.delta", ToolName: "secret tool", Status: "partial answer that must not leak"},
		},
		err: fmt.Errorf("provider secret response: %w", service.ErrChatUpstreamProtocol),
	}
	handler := newTestDashScopeChatHandler(t, &fakeChatUseCase{turn: turn}, nil, nil)
	request := authenticatedDashScopeRequest(http.MethodPost, `{"input":{"prompt":"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)
	}
	for _, forbidden := range []string{`"finish_reason":"stop"`, "partial answer", "provider secret response", "secret tool"} {
		if strings.Contains(body, forbidden) {
			t.Fatalf("failure stream exposed %q: %s", forbidden, body)
		}
	}
}

func TestDashScopeChatPreparationBusyReturnsRetryAfter(t *testing.T) {
	handler := newTestDashScopeChatHandler(t, &fakeChatUseCase{err: service.ErrChatConversationBusy}, nil, nil)
	request := authenticatedDashScopeRequest(http.MethodPost, `{"input":{"prompt":"hello","session_id":"conv_0123456789abcdef01234567"}}`)
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusConflict || response.Header().Get("Retry-After") != "1" || !strings.Contains(response.Body.String(), `"code":"CHAT_CONVERSATION_BUSY"`) {
		t.Fatalf("status=%d headers=%#v body=%s", response.Code, response.Header(), response.Body.String())
	}
}

func newTestDashScopeChatHandler(t *testing.T, chat ChatUseCase, origins []string, logger *log.Logger) *DashScopeChatHandler {
	t.Helper()
	return newTestDashScopeChatHandlerWithLimit(t, chat, 4096, origins, logger)
}

func newTestDashScopeChatHandlerWithLimit(t *testing.T, chat ChatUseCase, maxBody int64, origins []string, logger *log.Logger) *DashScopeChatHandler {
	t.Helper()
	handler, err := NewDashScopeChatHandler(chat, DashScopeChatOptions{
		AppID:          testDashScopeAppID,
		AuthToken:      testChatToken,
		AllowedOrigins: origins,
		MaxBodyBytes:   maxBody,
		RunTimeout:     time.Minute,
		Logger:         logger,
	})
	if err != nil {
		t.Fatalf("NewDashScopeChatHandler() error = %v", err)
	}
	return handler
}

func authenticatedDashScopeRequest(method, body string) *http.Request {
	request := httptest.NewRequest(method, "/api/v1/apps/"+testDashScopeAppID+"/completion", strings.NewReader(body))
	request.SetPathValue("app_id", testDashScopeAppID)
	request.Header.Set("xtoken", testChatToken)
	request.Header.Set("Content-Type", "application/json")
	return request
}

func dashScopeSSEBlocks(body string) []string {
	trimmed := strings.TrimSpace(body)
	if trimmed == "" {
		return nil
	}
	return strings.Split(trimmed, "\n\n")
}

func decodeDashScopeSSEData(t *testing.T, block string) dashScopeResultPayload {
	t.Helper()
	for _, line := range strings.Split(block, "\n") {
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		var payload dashScopeResultPayload
		if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &payload); err != nil {
			t.Fatalf("decode SSE data: %v", err)
		}
		return payload
	}
	t.Fatalf("SSE block has no data: %s", block)
	return dashScopeResultPayload{}
}