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

package handler

import (
	"bytes"
	"context"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

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

// DashScopeChatOptions configures the public DashScope-compatible chat
// transport. Its token authenticates callers of this service and is never a
// provider credential.
type DashScopeChatOptions struct {
	AppID                 string
	AuthToken             string
	AllowLegacyShortToken bool
	AllowedOrigins        []string
	MaxBodyBytes          int64
	RunTimeout            time.Duration
	Logger                *log.Logger
}

// DashScopeChatHandler adapts the provider-neutral ChatUseCase to the narrow
// DashScope application completion contract used by existing callers.
type DashScopeChatHandler struct {
	chat           ChatUseCase
	appID          string
	authToken      string
	allowedOrigins map[string]struct{}
	maxBodyBytes   int64
	runTimeout     time.Duration
	logger         *log.Logger
}

// NewDashScopeChatHandler constructs the protected compatibility transport.
func NewDashScopeChatHandler(chat ChatUseCase, options DashScopeChatOptions) (*DashScopeChatHandler, error) {
	if chat == nil {
		return nil, errors.New("chat service is required")
	}
	if !validDashScopeAppID(options.AppID) {
		return nil, errors.New("DashScope-compatible app ID is invalid")
	}
	if !validChatToken(options.AuthToken, options.AllowLegacyShortToken) {
		return nil, errors.New("DashScope-compatible auth token is outside the configured printable ASCII length policy")
	}
	if options.MaxBodyBytes <= 0 || options.MaxBodyBytes > maximumChatBodyBytes {
		return nil, errors.New("DashScope-compatible max body bytes is outside the supported range")
	}
	if options.RunTimeout <= 0 || options.RunTimeout > maximumChatRunTime {
		return nil, errors.New("DashScope-compatible run timeout is outside the supported range")
	}
	allowedOrigins := make(map[string]struct{}, len(options.AllowedOrigins))
	for _, origin := range options.AllowedOrigins {
		normalized, ok := validChatOrigin(origin)
		if !ok || normalized != origin {
			return nil, errors.New("DashScope-compatible allowed origins must be normalized exact HTTP(S) origins")
		}
		allowedOrigins[origin] = struct{}{}
	}
	logger := options.Logger
	if logger == nil {
		logger = log.New(io.Discard, "", 0)
	}
	return &DashScopeChatHandler{
		chat:           chat,
		appID:          options.AppID,
		authToken:      options.AuthToken,
		allowedOrigins: allowedOrigins,
		maxBodyBytes:   options.MaxBodyBytes,
		runTimeout:     options.RunTimeout,
		logger:         logger,
	}, nil
}

// ServeHTTP handles compatible application completion requests and CORS
// preflight. It buffers the upstream answer until strict ChatUseCase success.
func (h *DashScopeChatHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	started := time.Now()
	requestID := safeRequestID(r.Header.Get("X-Request-ID"))
	w.Header().Set("X-Request-ID", requestID)
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("X-Content-Type-Options", "nosniff")

	if r.PathValue("app_id") != h.appID {
		h.writeError(w, http.StatusNotFound, requestID, "CHAT_APP_NOT_FOUND", "Chat application was not found.")
		h.logResult(requestID, "app_not_found", false, started)
		return
	}
	if r.Method == http.MethodOptions {
		h.handlePreflight(w, r, requestID, started)
		return
	}
	if r.Method != http.MethodPost {
		w.Header().Set("Allow", http.MethodPost+", "+http.MethodOptions)
		h.writeError(w, http.StatusMethodNotAllowed, requestID, "CHAT_METHOD_NOT_ALLOWED", "Chat endpoint accepts POST requests only.")
		h.logResult(requestID, "method_not_allowed", false, started)
		return
	}
	if !h.authorizeOrigin(w, r.Header.Get("Origin")) {
		h.writeError(w, http.StatusForbidden, requestID, "CHAT_ORIGIN_FORBIDDEN", "Chat browser origin is not allowed.")
		h.logResult(requestID, "forbidden_origin", false, started)
		return
	}
	if !h.validToken(r.Header.Get("xtoken")) {
		w.Header().Set("WWW-Authenticate", `XToken realm="fire-safety-ymd-chat"`)
		h.writeError(w, http.StatusUnauthorized, requestID, "CHAT_AUTH_INVALID", "Chat authentication failed.")
		h.logResult(requestID, "unauthorized", false, started)
		return
	}
	if !isJSONContentType(r.Header.Get("Content-Type")) {
		h.writeError(w, http.StatusUnsupportedMediaType, requestID, "CHAT_CONTENT_TYPE_INVALID", "Content-Type must be application/json.")
		h.logResult(requestID, "unsupported_media_type", false, started)
		return
	}

	body, tooLarge, err := readBoundedBody(r.Body, h.maxBodyBytes)
	if err != nil {
		h.writeError(w, http.StatusBadRequest, requestID, "CHAT_REQUEST_INVALID", "Unable to read chat request.")
		h.logResult(requestID, "read_error", false, started)
		return
	}
	if tooLarge {
		h.writeError(w, http.StatusRequestEntityTooLarge, requestID, "CHAT_REQUEST_BODY_TOO_LARGE", "Chat request body exceeds the configured limit.")
		h.logResult(requestID, "body_too_large", false, started)
		return
	}
	request, err := decodeDashScopeChatRequest(body)
	if err != nil {
		h.writeError(w, http.StatusBadRequest, requestID, "CHAT_REQUEST_INVALID", "Chat request JSON is invalid.")
		h.logResult(requestID, "invalid_request", false, started)
		return
	}

	runCtx, cancelRun := context.WithTimeout(r.Context(), h.runTimeout)
	defer cancelRun()
	turn, err := h.chat.Prepare(runCtx, service.ChatRequest{
		Message:        request.Prompt,
		ConversationID: request.SessionID,
		RequestID:      requestID,
	})
	if err != nil {
		status, code, message := publicChatError(err)
		if errors.Is(err, service.ErrChatConversationBusy) {
			w.Header().Set("Retry-After", "1")
		}
		if errors.Is(err, service.ErrChatCapacityReached) {
			w.Header().Set("Retry-After", "30")
		}
		h.writeError(w, status, requestID, code, message)
		h.logResult(requestID, code, request.SessionID != "", started)
		return
	}
	defer turn.Close()

	flusher, ok := w.(http.Flusher)
	if !ok {
		h.writeError(w, http.StatusInternalServerError, requestID, "CHAT_STREAM_UNSUPPORTED", "Chat streaming is unavailable.")
		h.logResult(requestID, "stream_unsupported", turn.Reused(), started)
		return
	}
	if deadline, ok := runCtx.Deadline(); ok {
		_ = http.NewResponseController(w).SetWriteDeadline(deadline.Add(5 * time.Second))
	}
	w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	w.Header().Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)
	if err := writeDashScopeSSE(w, flusher, 1, "result", dashScopeResultPayload{
		Output:    dashScopeOutput{SessionID: turn.ConversationID(), FinishReason: "null"},
		Usage:     dashScopeUsage{},
		RequestID: requestID,
	}); err != nil {
		h.logResult(requestID, "client_write_failed", turn.Reused(), started)
		return
	}

	streamCtx, cancelStream := context.WithCancel(runCtx)
	defer cancelStream()
	type streamOutcome struct {
		result service.ChatResult
		err    error
	}
	outcome := make(chan streamOutcome, 1)
	go func() {
		result, streamErr := turn.Stream(streamCtx, nil)
		outcome <- streamOutcome{result: result, err: streamErr}
	}()

	heartbeat := time.NewTicker(chatHeartbeatInterval)
	defer heartbeat.Stop()
	contextDone := runCtx.Done()
	var result service.ChatResult
	var streamErr error
streamLoop:
	for {
		select {
		case completed := <-outcome:
			result = completed.result
			streamErr = completed.err
			break streamLoop
		case <-heartbeat.C:
			if err := writeChatHeartbeat(w, flusher); err != nil {
				cancelStream()
				h.logResult(requestID, "client_write_failed", turn.Reused(), started)
				return
			}
		case <-contextDone:
			cancelStream()
			contextDone = nil
		}
	}
	if streamErr != nil {
		_, code, message := publicChatError(streamErr)
		_ = writeDashScopeSSE(w, flusher, 2, "error", dashScopeStreamError{
			Code:      code,
			Message:   message,
			RequestID: requestID,
			SessionID: turn.ConversationID(),
		})
		h.logResult(requestID, code, turn.Reused(), started)
		return
	}
	if err := writeDashScopeSSE(w, flusher, 2, "result", dashScopeResultPayload{
		Output: dashScopeOutput{
			SessionID:    result.ConversationID,
			FinishReason: "stop",
			Text:         result.Answer,
		},
		Usage: dashScopeUsage{Models: []dashScopeModelUsage{{
			InputTokens:  result.Usage.Input,
			OutputTokens: result.Usage.Output,
			ModelID:      result.ModelID,
		}}},
		RequestID: requestID,
	}); err != nil {
		h.logResult(requestID, "client_write_failed", turn.Reused(), started)
		return
	}
	h.logResult(requestID, "success", turn.Reused(), started)
}

type dashScopeChatRequest struct {
	Prompt    string
	SessionID string
}

func decodeDashScopeChatRequest(body []byte) (dashScopeChatRequest, error) {
	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 || trimmed[0] != '{' {
		return dashScopeChatRequest{}, errors.New("request must be an object")
	}
	var envelope struct {
		Input      json.RawMessage `json:"input"`
		Parameters json.RawMessage `json:"parameters"`
		Debug      json.RawMessage `json:"debug"`
	}
	decoder := json.NewDecoder(bytes.NewReader(trimmed))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&envelope); err != nil {
		return dashScopeChatRequest{}, errors.New("request is invalid")
	}
	if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
		return dashScopeChatRequest{}, errors.New("request contains trailing JSON")
	}

	input, err := decodeDashScopeInput(envelope.Input)
	if err != nil {
		return dashScopeChatRequest{}, err
	}
	if len(envelope.Parameters) > 0 {
		var parameters struct {
			IncrementalOutput json.RawMessage `json:"incremental_output"`
		}
		if err := decodeStrictObject(envelope.Parameters, &parameters); err != nil {
			return dashScopeChatRequest{}, errors.New("parameters must be a supported object")
		}
		if len(parameters.IncrementalOutput) > 0 {
			var incrementalOutput bool
			if err := json.Unmarshal(parameters.IncrementalOutput, &incrementalOutput); err != nil || bytes.Equal(bytes.TrimSpace(parameters.IncrementalOutput), []byte("null")) {
				return dashScopeChatRequest{}, errors.New("parameters.incremental_output must be a boolean")
			}
		}
	}
	if len(envelope.Debug) > 0 {
		var debug map[string]json.RawMessage
		if err := decodeStrictObject(envelope.Debug, &debug); err != nil || len(debug) != 0 {
			return dashScopeChatRequest{}, errors.New("debug must be an empty object")
		}
	}
	return input, nil
}

func decodeDashScopeInput(raw json.RawMessage) (dashScopeChatRequest, error) {
	var input struct {
		Prompt    *string         `json:"prompt"`
		SessionID json.RawMessage `json:"session_id"`
	}
	if err := decodeStrictObject(raw, &input); err != nil || input.Prompt == nil || strings.TrimSpace(*input.Prompt) == "" {
		return dashScopeChatRequest{}, errors.New("input.prompt is required")
	}
	request := dashScopeChatRequest{Prompt: *input.Prompt}
	if len(input.SessionID) > 0 {
		if err := json.Unmarshal(input.SessionID, &request.SessionID); err != nil || strings.TrimSpace(request.SessionID) == "" {
			return dashScopeChatRequest{}, errors.New("input.session_id must be a non-empty string when provided")
		}
	}
	return request, nil
}

type dashScopeResultPayload struct {
	Output    dashScopeOutput `json:"output"`
	Usage     dashScopeUsage  `json:"usage"`
	RequestID string          `json:"request_id"`
}

type dashScopeOutput struct {
	SessionID    string `json:"session_id"`
	FinishReason string `json:"finish_reason"`
	Text         string `json:"text,omitempty"`
}

type dashScopeUsage struct {
	Models []dashScopeModelUsage `json:"models,omitempty"`
}

type dashScopeModelUsage struct {
	InputTokens  int64  `json:"input_tokens"`
	OutputTokens int64  `json:"output_tokens"`
	ModelID      string `json:"model_id,omitempty"`
}

type dashScopeStreamError struct {
	Code      string `json:"code"`
	Message   string `json:"message"`
	RequestID string `json:"request_id"`
	SessionID string `json:"session_id"`
}

func (h *DashScopeChatHandler) handlePreflight(w http.ResponseWriter, r *http.Request, requestID string, started time.Time) {
	origin := r.Header.Get("Origin")
	if origin == "" || !h.authorizeOrigin(w, origin) || r.Header.Get("Access-Control-Request-Method") != http.MethodPost ||
		!validDashScopePreflightHeaders(r.Header.Get("Access-Control-Request-Headers")) {
		h.writeError(w, http.StatusForbidden, requestID, "CHAT_ORIGIN_FORBIDDEN", "Chat browser origin or preflight request is not allowed.")
		h.logResult(requestID, "preflight_forbidden", false, started)
		return
	}
	w.Header().Set("Access-Control-Allow-Methods", http.MethodPost)
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type, xtoken, X-DashScope-SSE, X-Request-ID")
	w.Header().Set("Access-Control-Max-Age", "600")
	w.WriteHeader(http.StatusNoContent)
	h.logResult(requestID, "preflight_success", false, started)
}

func (h *DashScopeChatHandler) authorizeOrigin(w http.ResponseWriter, origin string) bool {
	if origin == "" {
		return true
	}
	if _, allowed := h.allowedOrigins[origin]; !allowed {
		return false
	}
	w.Header().Add("Vary", "Origin")
	w.Header().Set("Access-Control-Allow-Origin", origin)
	w.Header().Set("Access-Control-Expose-Headers", "X-Request-ID")
	return true
}

func (h *DashScopeChatHandler) validToken(value string) bool {
	if len(value) != len(h.authToken) {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(value), []byte(h.authToken)) == 1
}

func (h *DashScopeChatHandler) writeError(w http.ResponseWriter, status int, requestID, code, message string) {
	writeJSON(w, status, map[string]string{
		"code":       code,
		"message":    message,
		"request_id": requestID,
	})
}

func (h *DashScopeChatHandler) logResult(requestID, result string, reused bool, started time.Time) {
	h.logger.Printf("dashscope_chat_request request_id=%s result=%s reused=%t duration_ms=%d", requestID, result, reused, time.Since(started).Milliseconds())
}

func writeDashScopeSSE(w io.Writer, flusher http.Flusher, id int, event string, payload any) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	if _, err := fmt.Fprintf(w, "id: %d\nevent: %s\n:HTTP_STATUS/200\ndata: %s\n\n", id, event, data); err != nil {
		return err
	}
	flusher.Flush()
	return nil
}

func validDashScopePreflightHeaders(value string) bool {
	for _, header := range strings.Split(value, ",") {
		switch strings.ToLower(strings.TrimSpace(header)) {
		case "", "content-type", "xtoken", "x-dashscope-sse", "x-request-id":
		default:
			return false
		}
	}
	return true
}

func validDashScopeAppID(value string) bool {
	if value == "" || len(value) > 128 || strings.TrimSpace(value) != value {
		return false
	}
	for _, character := range value {
		if !(character >= 'a' && character <= 'z') &&
			!(character >= 'A' && character <= 'Z') &&
			!(character >= '0' && character <= '9') &&
			!strings.ContainsRune("_-", character) {
			return false
		}
	}
	return true
}