feat(chat): http handler + DTO + admin guard for endpoints/models/usage

2026-05-04 17:59:22 +08:00
parent 8b5dc70146
commit 175e9be36a
3 changed files with 1463 additions and 0 deletions
@@ -0,0 +1,390 @@
+package chat
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
+	"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
+)
+
+// ===== Fake service implementations for handler tests =====
+
+// fakeSessionResolver always resolves to fakeUID when a non-empty token is provided.
+type fakeSessionResolver struct {
+	uid uuid.UUID
+}
+
+func (f *fakeSessionResolver) ResolveSession(_ context.Context, token string) (uuid.UUID, bool, error) {
+	if token == "" {
+		return uuid.Nil, false, nil
+	}
+	return f.uid, true, nil
+}
+
+// fakeAdminLookup allows controlling IsAdmin per-user.
+type fakeAdminLookup struct {
+	adminIDs map[uuid.UUID]bool
+}
+
+func newFakeAdminLookup(admins ...uuid.UUID) *fakeAdminLookup {
+	m := make(map[uuid.UUID]bool, len(admins))
+	for _, id := range admins {
+		m[id] = true
+	}
+	return &fakeAdminLookup{adminIDs: m}
+}
+
+func (f *fakeAdminLookup) IsAdmin(_ context.Context, id uuid.UUID) (bool, error) {
+	return f.adminIDs[id], nil
+}
+
+// ===== Fake ConversationService =====
+
+type fakeConvSvc struct {
+	createFunc func(ctx context.Context, userID uuid.UUID, in CreateConversationInput) (*Conversation, error)
+	listFunc   func(ctx context.Context, userID uuid.UUID, f ConversationFilter) ([]Conversation, error)
+	getFunc    func(ctx context.Context, userID, id uuid.UUID) (*Conversation, []Message, error)
+}
+
+func (f *fakeConvSvc) Create(ctx context.Context, userID uuid.UUID, in CreateConversationInput) (*Conversation, error) {
+	return f.createFunc(ctx, userID, in)
+}
+func (f *fakeConvSvc) List(ctx context.Context, userID uuid.UUID, filter ConversationFilter) ([]Conversation, error) {
+	if f.listFunc != nil {
+		return f.listFunc(ctx, userID, filter)
+	}
+	return nil, nil
+}
+func (f *fakeConvSvc) Get(ctx context.Context, userID, id uuid.UUID) (*Conversation, []Message, error) {
+	if f.getFunc != nil {
+		return f.getFunc(ctx, userID, id)
+	}
+	return nil, nil, errs.New(errs.CodeChatConversationNotFound, "not found")
+}
+func (f *fakeConvSvc) UpdateTitle(_ context.Context, _, _ uuid.UUID, _ string) error { return nil }
+func (f *fakeConvSvc) SoftDelete(_ context.Context, _, _ uuid.UUID) error            { return nil }
+
+// ===== Fake MessageService =====
+
+type fakeMsgSvc struct {
+	sendFunc func(ctx context.Context, userID, convID uuid.UUID, content string, ids []uuid.UUID) (int64, int64, error)
+	retryErr error
+}
+
+func (f *fakeMsgSvc) Send(ctx context.Context, userID, convID uuid.UUID, content string, ids []uuid.UUID) (int64, int64, error) {
+	if f.sendFunc != nil {
+		return f.sendFunc(ctx, userID, convID, content, ids)
+	}
+	return 1, 2, nil
+}
+func (f *fakeMsgSvc) Stream(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ int64) (<-chan SSEEvent, error) {
+	ch := make(chan SSEEvent)
+	close(ch)
+	return ch, nil
+}
+func (f *fakeMsgSvc) Cancel(_ context.Context, _ uuid.UUID, _ int64) error { return nil }
+func (f *fakeMsgSvc) Retry(_ context.Context, _ uuid.UUID, msgID int64) (int64, error) {
+	if f.retryErr != nil {
+		return 0, f.retryErr
+	}
+	return msgID + 100, nil
+}
+func (f *fakeMsgSvc) ListMessages(_ context.Context, _, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
+	return nil, nil
+}
+
+// ===== Fake TemplateService =====
+
+type fakeTplSvc struct {
+	createFunc func(ctx context.Context, userID uuid.UUID, name, content string, scope TemplateScope) (*PromptTemplate, error)
+}
+
+func (f *fakeTplSvc) List(_ context.Context, _ uuid.UUID) ([]PromptTemplate, error) { return nil, nil }
+func (f *fakeTplSvc) Create(ctx context.Context, userID uuid.UUID, name, content string, scope TemplateScope) (*PromptTemplate, error) {
+	if f.createFunc != nil {
+		return f.createFunc(ctx, userID, name, content, scope)
+	}
+	return &PromptTemplate{ID: uuid.New(), Name: name, Content: content, Scope: scope, CreatedAt: time.Now(), UpdatedAt: time.Now()}, nil
+}
+func (f *fakeTplSvc) Update(_ context.Context, _, _ uuid.UUID, _, _ string) error { return nil }
+func (f *fakeTplSvc) Delete(_ context.Context, _, _ uuid.UUID) error              { return nil }
+
+// ===== Fake EndpointService =====
+
+type fakeEpSvc struct {
+	listEndpointsErr error
+}
+
+func (f *fakeEpSvc) ListEndpoints(_ context.Context) ([]LLMEndpoint, error) {
+	if f.listEndpointsErr != nil {
+		return nil, f.listEndpointsErr
+	}
+	return nil, nil
+}
+func (f *fakeEpSvc) CreateEndpoint(_ context.Context, _ uuid.UUID, in CreateEndpointInput) (*LLMEndpoint, error) {
+	return &LLMEndpoint{ID: uuid.New(), Provider: in.Provider, DisplayName: in.DisplayName, BaseURL: in.BaseURL, CreatedAt: time.Now(), UpdatedAt: time.Now()}, nil
+}
+func (f *fakeEpSvc) UpdateEndpoint(_ context.Context, _ uuid.UUID, _ UpdateEndpointInput) error {
+	return nil
+}
+func (f *fakeEpSvc) DeleteEndpoint(_ context.Context, _ uuid.UUID) error { return nil }
+func (f *fakeEpSvc) TestEndpoint(_ context.Context, _ uuid.UUID) error   { return nil }
+func (f *fakeEpSvc) ListModels(_ context.Context, _ bool) ([]LLMModel, error) {
+	return nil, nil
+}
+func (f *fakeEpSvc) CreateModel(_ context.Context, in CreateModelInput) (*LLMModel, error) {
+	return &LLMModel{ID: uuid.New(), ModelID: in.ModelID, DisplayName: in.DisplayName, CreatedAt: time.Now(), UpdatedAt: time.Now()}, nil
+}
+func (f *fakeEpSvc) UpdateModel(_ context.Context, _ uuid.UUID, _ UpdateModelInput) error {
+	return nil
+}
+func (f *fakeEpSvc) DeleteModel(_ context.Context, _ uuid.UUID) error { return nil }
+
+// ===== Fake UsageService =====
+
+type fakeUsageSvc struct{}
+
+func (f *fakeUsageSvc) Insert(_ context.Context, _ UsageRecord) error { return nil }
+func (f *fakeUsageSvc) SummarizeByUser(_ context.Context, _, _ time.Time) ([]UserUsageRow, error) {
+	return nil, nil
+}
+func (f *fakeUsageSvc) SummarizeByModel(_ context.Context, _, _ time.Time) ([]ModelUsageRow, error) {
+	return nil, nil
+}
+func (f *fakeUsageSvc) SummarizeByDay(_ context.Context, _, _ time.Time) ([]DayUsageRow, error) {
+	return nil, nil
+}
+func (f *fakeUsageSvc) UserDaily(_ context.Context, _ uuid.UUID, _, _ time.Time) ([]DayUsageRow, error) {
+	return nil, nil
+}
+
+// ===== Test helper =====
+
+// newTestHandler builds a Handler with the given services and mounts it on a chi router.
+// The resolver resolves any non-empty token to uid; adminIDs are treated as admins.
+func newTestHandler(
+	convSvc ConversationService,
+	msgSvc MessageService,
+	tplSvc TemplateService,
+	epSvc EndpointService,
+	uid uuid.UUID,
+	adminLookup AdminLookup,
+) (*Handler, chi.Router) {
+	resolver := &fakeSessionResolver{uid: uid}
+	upload := NewUploader(newFakeRepo(), nil, nil, 1<<20)
+	h := NewHandler(convSvc, msgSvc, tplSvc, epSvc, &fakeUsageSvc{}, upload, resolver, adminLookup)
+	r := chi.NewRouter()
+	r.Use(middleware.RequestID)
+	h.Mount(r)
+	return h, r
+}
+
+// authRequest creates a request with a Bearer token so Auth middleware passes.
+func authRequest(method, path string, body []byte) *http.Request {
+	var r *http.Request
+	if body != nil {
+		r = httptest.NewRequest(method, path, bytes.NewReader(body))
+		r.Header.Set("Content-Type", "application/json")
+	} else {
+		r = httptest.NewRequest(method, path, nil)
+	}
+	r.Header.Set("Authorization", "Bearer test-token")
+	return r
+}
+
+// ===== Tests =====
+
+// TestHandler_SendMessage_Returns200WithStreamURL verifies that sendMessage
+// calls the service with the right arguments and returns a SendMessageResp
+// with a correctly formed StreamURL.
+func TestHandler_SendMessage_Returns200WithStreamURL(t *testing.T) {
+	uid := uuid.New()
+	convID := uuid.New()
+
+	var capturedConvID uuid.UUID
+	var capturedContent string
+	msgSvc := &fakeMsgSvc{
+		sendFunc: func(_ context.Context, _ uuid.UUID, cID uuid.UUID, content string, _ []uuid.UUID) (int64, int64, error) {
+			capturedConvID = cID
+			capturedContent = content
+			return 10, 20, nil
+		},
+	}
+
+	_, r := newTestHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
+
+	body, _ := json.Marshal(SendMessageReq{Content: "hello"})
+	req := authRequest(http.MethodPost, "/api/v1/conversations/"+convID.String()+"/messages", body)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	require.Equal(t, http.StatusOK, w.Code)
+	var resp SendMessageResp
+	require.NoError(t, json.NewDecoder(w.Body).Decode(&resp))
+	assert.Equal(t, int64(10), resp.UserMessageID)
+	assert.Equal(t, int64(20), resp.AssistantMessageID)
+	assert.Contains(t, resp.StreamURL, convID.String())
+	assert.Equal(t, convID, capturedConvID)
+	assert.Equal(t, "hello", capturedContent)
+}
+
+// TestHandler_CreateConversation_LocksTemplateContent verifies that the handler
+// passes the TemplateID and SystemPrompt fields through to the service unchanged.
+// The service is responsible for resolving template content; the handler must
+// not override or ignore these fields.
+func TestHandler_CreateConversation_LocksTemplateContent(t *testing.T) {
+	uid := uuid.New()
+	templateID := uuid.New()
+	customPrompt := "my custom system prompt"
+
+	var capturedInput CreateConversationInput
+	convSvc := &fakeConvSvc{
+		createFunc: func(_ context.Context, _ uuid.UUID, in CreateConversationInput) (*Conversation, error) {
+			capturedInput = in
+			return &Conversation{
+				ID:        uuid.New(),
+				UserID:    uid,
+				ModelID:   in.ModelID,
+				CreatedAt: time.Now(),
+				UpdatedAt: time.Now(),
+			}, nil
+		},
+	}
+
+	_, r := newTestHandler(convSvc, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
+
+	modelID := uuid.New()
+	body, _ := json.Marshal(CreateConversationReq{
+		ModelID:      modelID,
+		TemplateID:   &templateID,
+		SystemPrompt: &customPrompt,
+	})
+	req := authRequest(http.MethodPost, "/api/v1/conversations", body)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	require.Equal(t, http.StatusCreated, w.Code)
+	// Handler must forward both TemplateID and SystemPrompt to the service
+	// (service decides precedence, not handler).
+	assert.Equal(t, modelID, capturedInput.ModelID)
+	require.NotNil(t, capturedInput.TemplateID)
+	assert.Equal(t, templateID, *capturedInput.TemplateID)
+	require.NotNil(t, capturedInput.SystemPrompt)
+	assert.Equal(t, customPrompt, *capturedInput.SystemPrompt)
+}
+
+// TestHandler_AdminEndpoint_403WhenNotAdmin verifies that a non-admin user
+// receives 403 when calling an admin-only endpoint.
+func TestHandler_AdminEndpoint_403WhenNotAdmin(t *testing.T) {
+	uid := uuid.New()
+	// uid is NOT in the admin lookup.
+	_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
+
+	req := authRequest(http.MethodGet, "/api/v1/admin/llm-endpoints", nil)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusForbidden, w.Code)
+	var body map[string]any
+	require.NoError(t, json.NewDecoder(w.Body).Decode(&body))
+	assert.Equal(t, string(errs.CodeForbidden), body["code"])
+}
+
+// TestHandler_RetryMessage_404OnUnknownID verifies that when the message
+// service returns a not-found error, the handler maps it to 404.
+func TestHandler_RetryMessage_404OnUnknownID(t *testing.T) {
+	uid := uuid.New()
+	msgSvc := &fakeMsgSvc{
+		retryErr: errs.New(errs.CodeChatMessageNotFound, "message not found"),
+	}
+
+	_, r := newTestHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
+
+	req := authRequest(http.MethodPost, "/api/v1/messages/99999/retry", nil)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusNotFound, w.Code)
+	var body map[string]any
+	require.NoError(t, json.NewDecoder(w.Body).Decode(&body))
+	assert.Equal(t, string(errs.CodeChatMessageNotFound), body["code"])
+}
+
+// TestHandler_CreateTemplate_SystemScopeRequiresAdmin verifies that a
+// non-admin user cannot create a system-scoped template and receives 403.
+func TestHandler_CreateTemplate_SystemScopeRequiresAdmin(t *testing.T) {
+	uid := uuid.New()
+	// uid is not an admin.
+	_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
+
+	body, _ := json.Marshal(map[string]string{
+		"name":    "sys tpl",
+		"content": "You are a system assistant",
+		"scope":   "system",
+	})
+	req := authRequest(http.MethodPost, "/api/v1/prompt-templates", body)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusForbidden, w.Code)
+}
+
+// TestHandler_CreateTemplate_SystemScopeAllowedForAdmin verifies that an
+// admin user can create a system-scoped template.
+func TestHandler_CreateTemplate_SystemScopeAllowedForAdmin(t *testing.T) {
+	uid := uuid.New()
+	// uid is an admin.
+	_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup(uid))
+
+	body, _ := json.Marshal(map[string]string{
+		"name":    "sys tpl",
+		"content": "You are a system assistant",
+		"scope":   "system",
+	})
+	req := authRequest(http.MethodPost, "/api/v1/prompt-templates", body)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusCreated, w.Code)
+}
+
+// TestHandler_Unauthenticated_Returns401 verifies that requests without a
+// Bearer token are rejected by the Auth middleware with 401.
+func TestHandler_Unauthenticated_Returns401(t *testing.T) {
+	uid := uuid.New()
+	_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
+
+	req := httptest.NewRequest(http.MethodGet, "/api/v1/conversations", nil)
+	// No Authorization header.
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusUnauthorized, w.Code)
+}
+
+// TestHandler_AdminAllowed_ListEndpoints verifies that an admin can call a
+// admin-only endpoint and receives a 200.
+func TestHandler_AdminAllowed_ListEndpoints(t *testing.T) {
+	uid := uuid.New()
+	_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup(uid))
+
+	req := authRequest(http.MethodGet, "/api/v1/admin/llm-endpoints", nil)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	// Response should be a JSON array (empty is fine).
+	assert.True(t, strings.HasPrefix(w.Body.String(), "[") || w.Body.String() == "[]\n")
+}