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

package handler

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

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

const (
	maximumChatBodyBytes  = int64(1024 * 1024)
	maximumChatRunTime    = 30 * time.Minute
	chatHeartbeatInterval = 15 * time.Second
)

// ChatUseCase is the inbound chat capability consumed by ChatHandler.
type ChatUseCase interface {
	Prepare(context.Context, service.ChatRequest) (service.ChatTurn, error)
}

// ChatOptions contains the independently authenticated user-facing transport settings.
type ChatOptions struct {
	AuthToken             string
	AllowLegacyShortToken bool
	AllowedOrigins        []string
	MaxBodyBytes          int64
	RunTimeout            time.Duration
	Logger                *log.Logger
}

// ChatHandler exposes a bounded SSE chat API without exposing provider credentials or sessions.
type ChatHandler struct {
	chat           ChatUseCase
	authToken      string
	allowedOrigins map[string]struct{}
	maxBodyBytes   int64
	runTimeout     time.Duration
	logger         *log.Logger
}

// NewChatHandler constructs the protected chat transport.
func NewChatHandler(chat ChatUseCase, options ChatOptions) (*ChatHandler, error) {
	if chat == nil {
		return nil, errors.New("chat service is required")
	}
	if !validChatToken(options.AuthToken, options.AllowLegacyShortToken) {
		return nil, errors.New("chat auth token is outside the configured printable ASCII length policy")
	}
	if options.MaxBodyBytes <= 0 || options.MaxBodyBytes > maximumChatBodyBytes {
		return nil, errors.New("chat max body bytes is outside the supported range")
	}
	if options.RunTimeout <= 0 || options.RunTimeout > maximumChatRunTime {
		return nil, errors.New("chat 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("chat 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 &ChatHandler{
		chat:           chat,
		authToken:      options.AuthToken,
		allowedOrigins: allowedOrigins,
		maxBodyBytes:   options.MaxBodyBytes,
		runTimeout:     options.RunTimeout,
		logger:         logger,
	}, nil
}

// ServeHTTP handles POST chat turns and browser CORS preflight.
func (h *ChatHandler) 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.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.validAuthorization(r.Header.Get("Authorization")) {
		w.Header().Set("WWW-Authenticate", `Bearer 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
	}
	if !acceptsEventStream(r.Header.Get("Accept")) {
		h.writeError(w, http.StatusNotAcceptable, requestID, "CHAT_ACCEPT_INVALID", "Accept must allow text/event-stream.")
		h.logResult(requestID, "not_acceptable", 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 := decodeChatRequest(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.Message,
		ConversationID: request.ConversationID,
		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.ConversationID != "", 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 := writeChatSSE(w, flusher, "conversation", map[string]any{
		"conversation_id": turn.ConversationID(),
		"reused":          turn.Reused(),
	}); 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
	}
	progress := make(chan service.AgentTraceEvent)
	outcome := make(chan streamOutcome, 1)
	go func() {
		result, err := turn.Stream(streamCtx, func(event service.AgentTraceEvent) {
			if !exposeProgressEvent(event.Event) {
				return
			}
			select {
			case progress <- event:
			case <-streamCtx.Done():
			}
		})
		outcome <- streamOutcome{result: result, err: err}
	}()

	heartbeat := time.NewTicker(chatHeartbeatInterval)
	defer heartbeat.Stop()
	contextDone := runCtx.Done()
	var result service.ChatResult
	var streamErr error
streamLoop:
	for {
		select {
		case event := <-progress:
			if err := writeChatSSE(w, flusher, "progress", event); err != nil {
				cancelStream()
				h.logResult(requestID, "client_write_failed", turn.Reused(), started)
				return
			}
		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)
		_ = writeChatSSE(w, flusher, "error", map[string]string{
			"code":            code,
			"message":         message,
			"conversation_id": turn.ConversationID(),
			"request_id":      requestID,
		})
		h.logResult(requestID, code, turn.Reused(), started)
		return
	}
	if err := writeChatSSE(w, flusher, "message", map[string]string{
		"conversation_id": result.ConversationID,
		"answer":          result.Answer,
	}); err != nil {
		h.logResult(requestID, "client_write_failed", turn.Reused(), started)
		return
	}
	if err := writeChatSSE(w, flusher, "done", map[string]any{
		"conversation_id": result.ConversationID,
		"run_id":          result.RunID,
		"usage":           result.Usage,
	}); err != nil {
		h.logResult(requestID, "client_write_failed", turn.Reused(), started)
		return
	}
	h.logResult(requestID, "success", turn.Reused(), started)
}

type chatRequestPayload struct {
	Message        string
	ConversationID string
}

func decodeChatRequest(body []byte) (chatRequestPayload, error) {
	var payload struct {
		Message        *string         `json:"message"`
		ConversationID json.RawMessage `json:"conversation_id"`
	}
	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 || trimmed[0] != '{' {
		return chatRequestPayload{}, errors.New("chat request must be an object")
	}
	decoder := json.NewDecoder(bytes.NewReader(trimmed))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&payload); err != nil || payload.Message == nil {
		return chatRequestPayload{}, errors.New("chat request is invalid")
	}
	if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
		return chatRequestPayload{}, errors.New("chat request contains trailing JSON")
	}
	request := chatRequestPayload{Message: *payload.Message}
	if len(payload.ConversationID) > 0 {
		if err := json.Unmarshal(payload.ConversationID, &request.ConversationID); err != nil || strings.TrimSpace(request.ConversationID) == "" {
			return chatRequestPayload{}, errors.New("conversation_id must be a non-empty string when provided")
		}
	}
	return request, nil
}

func (h *ChatHandler) 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 ||
		!validPreflightHeaders(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", "Authorization, Content-Type")
	w.Header().Set("Access-Control-Max-Age", "600")
	w.WriteHeader(http.StatusNoContent)
	h.logResult(requestID, "preflight_success", false, started)
}

func (h *ChatHandler) 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 *ChatHandler) validAuthorization(value string) bool {
	expected := "Bearer " + h.authToken
	if len(value) != len(expected) {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(value), []byte(expected)) == 1
}

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

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

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

func writeChatHeartbeat(w io.Writer, flusher http.Flusher) error {
	if _, err := io.WriteString(w, ": keepalive\n\n"); err != nil {
		return err
	}
	flusher.Flush()
	return nil
}

func publicChatError(err error) (int, string, string) {
	switch {
	case errors.Is(err, service.ErrChatInvalidArgument):
		return http.StatusBadRequest, "CHAT_REQUEST_INVALID", "Chat request is invalid."
	case errors.Is(err, service.ErrChatConversationNotFound):
		return http.StatusNotFound, "CHAT_CONVERSATION_NOT_FOUND", "Chat conversation was not found or has expired."
	case errors.Is(err, service.ErrChatConversationBusy):
		return http.StatusConflict, "CHAT_CONVERSATION_BUSY", "Chat conversation already has an active run."
	case errors.Is(err, service.ErrChatCapacityReached):
		return http.StatusServiceUnavailable, "CHAT_CAPACITY_REACHED", "Chat session capacity is temporarily exhausted."
	case errors.Is(err, context.DeadlineExceeded):
		return http.StatusGatewayTimeout, "CHAT_UPSTREAM_TIMEOUT", "Chat run did not complete within its time limit."
	case errors.Is(err, context.Canceled):
		return http.StatusRequestTimeout, "CHAT_REQUEST_CANCELED", "Chat request was canceled."
	case errors.Is(err, service.ErrChatRunFailed):
		return http.StatusBadGateway, "CHAT_RUN_FAILED", "SuperAgent reported a failed run."
	case errors.Is(err, service.ErrChatUpstreamProtocol):
		return http.StatusBadGateway, "CHAT_UPSTREAM_PROTOCOL_ERROR", "SuperAgent returned an incomplete or invalid response."
	case errors.Is(err, service.ErrChatUpstreamUnavailable):
		return http.StatusBadGateway, "CHAT_UPSTREAM_UNAVAILABLE", "SuperAgent is unavailable."
	default:
		return http.StatusInternalServerError, "CHAT_INTERNAL_ERROR", "Chat request failed."
	}
}

func validChatToken(value string, allowLegacyShortToken bool) bool {
	if value == "" || len(value) > 4096 || (!allowLegacyShortToken && len(value) < 32) {
		return false
	}
	for _, character := range value {
		if character <= 0x20 || character >= 0x7f {
			return false
		}
	}
	return true
}

func validChatOrigin(value string) (string, bool) {
	parsed, err := url.Parse(value)
	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") ||
		parsed.User != nil || parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" ||
		(parsed.Path != "" && parsed.Path != "/") || value == "*" {
		return "", false
	}
	return parsed.Scheme + "://" + parsed.Host, true
}

func acceptsEventStream(value string) bool {
	if strings.TrimSpace(value) == "" {
		return true
	}
	for _, candidate := range strings.Split(value, ",") {
		mediaType, parameters, err := mime.ParseMediaType(strings.TrimSpace(candidate))
		if err != nil {
			continue
		}
		if quality, exists := parameters["q"]; exists {
			parsed, parseErr := strconv.ParseFloat(quality, 64)
			if parseErr != nil || parsed <= 0 || parsed > 1 {
				continue
			}
		}
		if mediaType == "*/*" || mediaType == "text/*" || strings.EqualFold(mediaType, "text/event-stream") {
			return true
		}
	}
	return false
}

func validPreflightHeaders(value string) bool {
	for _, header := range strings.Split(value, ",") {
		header = strings.ToLower(strings.TrimSpace(header))
		if header == "" {
			continue
		}
		if header != "authorization" && header != "content-type" {
			return false
		}
	}
	return true
}

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