feat: add ai gateway local core flow

apps/api/internal/httpapi/core_flow_integration_test.go

@@ -0,0 +1,298 @@
+package httpapi
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"runtime"
+	"sort"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func TestCoreLocalFlow(t *testing.T) {
+	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
+	if databaseURL == "" {
+		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the PostgreSQL integration flow")
+	}
+	ctx := context.Background()
+	applyMigration(t, ctx, databaseURL)
+
+	db, err := store.Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect store: %v", err)
+	}
+	defer db.Close()
+
+	handler := NewServer(config.Config{
+		AppEnv:            "test",
+		HTTPAddr:          ":0",
+		DatabaseURL:       databaseURL,
+		IdentityMode:      "hybrid",
+		JWTSecret:         "test-secret",
+		CORSAllowedOrigin: "*",
+	}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	suffix := time.Now().UnixNano()
+	suffixText := strconv.FormatInt(suffix, 10)
+	username := "smoke_admin_" + suffixText
+	password := "password123"
+
+	var registerResponse struct {
+		AccessToken string `json:"accessToken"`
+	}
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
+		"username":   username,
+		"email":      username + "@example.com",
+		"password":   password,
+		"tenantKey":  "manual-" + suffixText,
+		"tenantName": "Manual Tenant",
+	}, http.StatusCreated, &registerResponse)
+	if registerResponse.AccessToken == "" {
+		t.Fatal("register did not return access token")
+	}
+
+	var duplicateResponse map[string]any
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
+		"username": username,
+		"email":    username + "@example.com",
+		"password": password,
+	}, http.StatusConflict, &duplicateResponse)
+	if errorBody, ok := duplicateResponse["error"].(map[string]any); !ok || errorBody["message"] != "user already exists" {
+		t.Fatalf("unexpected duplicate response: %+v", duplicateResponse)
+	}
+
+	var loginResponse struct {
+		AccessToken string `json:"accessToken"`
+	}
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
+		"account":  username,
+		"password": password,
+	}, http.StatusOK, &loginResponse)
+	if loginResponse.AccessToken == "" {
+		t.Fatal("login did not return access token")
+	}
+
+	var apiKeyResponse struct {
+		Secret string `json:"secret"`
+		APIKey struct {
+			ID        string `json:"id"`
+			KeyPrefix string `json:"keyPrefix"`
+			Status    string `json:"status"`
+		} `json:"apiKey"`
+	}
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", loginResponse.AccessToken, map[string]any{
+		"name":   "smoke key",
+		"scopes": []string{"chat", "image", "video"},
+	}, http.StatusCreated, &apiKeyResponse)
+	if !strings.HasPrefix(apiKeyResponse.Secret, "sk-gw-") || apiKeyResponse.APIKey.Status != "active" {
+		t.Fatalf("unexpected api key response: %+v", apiKeyResponse)
+	}
+
+	var me map[string]any
+	doJSON(t, server.URL, http.MethodGet, "/api/v1/me", apiKeyResponse.Secret, nil, http.StatusOK, &me)
+	if me["apiKeyId"] == "" {
+		t.Fatalf("api key auth did not expose apiKeyId: %+v", me)
+	}
+	if me["tenantKey"] != "default" {
+		t.Fatalf("register should ignore public tenant fields and use default tenant: %+v", me)
+	}
+
+	testPool, err := pgxpool.New(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect test pool: %v", err)
+	}
+	defer testPool.Close()
+
+	inviteCode := "INVITE-" + suffixText
+	if _, err := testPool.Exec(ctx, `
+INSERT INTO gateway_invitations (invite_code, max_uses, metadata)
+VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
+		t.Fatalf("insert invitation: %v", err)
+	}
+	invitedUsername := "smoke_invited_" + suffixText
+	var invitedRegisterResponse struct {
+		AccessToken string `json:"accessToken"`
+	}
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
+		"username":       invitedUsername,
+		"email":          invitedUsername + "@example.com",
+		"password":       password,
+		"tenantKey":      "manual-invited-" + suffixText,
+		"tenantName":     "Manual Invited Tenant",
+		"invitationCode": inviteCode,
+	}, http.StatusCreated, &invitedRegisterResponse)
+	var invitedMe map[string]any
+	doJSON(t, server.URL, http.MethodGet, "/api/v1/me", invitedRegisterResponse.AccessToken, nil, http.StatusOK, &invitedMe)
+	if invitedMe["tenantKey"] != "default" {
+		t.Fatalf("invitation should not change tenant context: %+v", invitedMe)
+	}
+	var usedCount int
+	if err := testPool.QueryRow(ctx, `SELECT used_count FROM gateway_invitations WHERE invite_code = $1`, inviteCode).Scan(&usedCount); err != nil {
+		t.Fatalf("read invitation used_count: %v", err)
+	}
+	if usedCount != 1 {
+		t.Fatalf("invitation used_count = %d, want 1", usedCount)
+	}
+	var userMetadata []byte
+	if err := testPool.QueryRow(ctx, `SELECT metadata FROM gateway_users WHERE username = $1`, invitedUsername).Scan(&userMetadata); err != nil {
+		t.Fatalf("read invited user metadata: %v", err)
+	}
+	var metadata map[string]any
+	if err := json.Unmarshal(userMetadata, &metadata); err != nil {
+		t.Fatalf("decode invited user metadata: %v", err)
+	}
+	registration, ok := metadata["registration"].(map[string]any)
+	if !ok || registration["invitationCode"] != inviteCode {
+		t.Fatalf("invitation relationship was not recorded: %+v", metadata)
+	}
+
+	var platform struct {
+		ID          string `json:"id"`
+		Provider    string `json:"provider"`
+		PlatformKey string `json:"platformKey"`
+		Status      string `json:"status"`
+	}
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/platforms", loginResponse.AccessToken, map[string]any{
+		"provider":    "openai",
+		"platformKey": "openai-smoke-" + suffixText,
+		"name":        "OpenAI Smoke",
+		"baseUrl":     "https://api.openai.com/v1",
+		"authType":    "bearer",
+		"credentials": map[string]any{"mode": "simulation"},
+		"config":      map[string]any{"testMode": true},
+	}, http.StatusCreated, &platform)
+	if platform.ID == "" || platform.Status != "enabled" {
+		t.Fatalf("unexpected platform response: %+v", platform)
+	}
+
+	var taskResponse struct {
+		Task struct {
+			ID      string         `json:"id"`
+			Status  string         `json:"status"`
+			RunMode string         `json:"runMode"`
+			Result  map[string]any `json:"result"`
+		} `json:"task"`
+	}
+	doJSON(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
+		"model":      "gpt-4o-mini",
+		"runMode":    "simulation",
+		"simulation": true,
+		"messages":   []map[string]any{{"role": "user", "content": "ping"}},
+	}, http.StatusAccepted, &taskResponse)
+	if taskResponse.Task.ID == "" || taskResponse.Task.Status != "succeeded" || taskResponse.Task.RunMode != "simulation" {
+		t.Fatalf("unexpected task response: %+v", taskResponse.Task)
+	}
+
+	var taskDetail struct {
+		ID     string         `json:"id"`
+		Status string         `json:"status"`
+		Result map[string]any `json:"result"`
+	}
+	doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+taskResponse.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &taskDetail)
+	if taskDetail.Status != "succeeded" || taskDetail.Result["id"] == "" {
+		t.Fatalf("unexpected task detail: %+v", taskDetail)
+	}
+
+	req, err := http.NewRequest(http.MethodGet, server.URL+"/api/v1/tasks/"+taskResponse.Task.ID+"/events", nil)
+	if err != nil {
+		t.Fatalf("build events request: %v", err)
+	}
+	req.Header.Set("Authorization", "Bearer "+apiKeyResponse.Secret)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("events request: %v", err)
+	}
+	defer resp.Body.Close()
+	body, _ := io.ReadAll(resp.Body)
+	if resp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("task.completed")) {
+		t.Fatalf("unexpected events response status=%d body=%s", resp.StatusCode, string(body))
+	}
+}
+
+func TestOriginAllowedSupportsCommaSeparatedOrigins(t *testing.T) {
+	allowed := "http://localhost:5178, http://127.0.0.1:5178"
+	if !originAllowed("http://localhost:5178", allowed) {
+		t.Fatal("localhost origin should be allowed")
+	}
+	if !originAllowed("http://127.0.0.1:5178", allowed) {
+		t.Fatal("127.0.0.1 origin should be allowed")
+	}
+	if originAllowed("http://127.0.0.1:5179", allowed) {
+		t.Fatal("unexpected origin should not be allowed")
+	}
+}
+
+func applyMigration(t *testing.T, ctx context.Context, databaseURL string) {
+	t.Helper()
+	_, filename, _, _ := runtime.Caller(0)
+	migrationFiles, err := filepath.Glob(filepath.Join(filepath.Dir(filename), "..", "..", "migrations", "*.sql"))
+	if err != nil {
+		t.Fatalf("read migration files: %v", err)
+	}
+	sort.Strings(migrationFiles)
+	pool, err := pgxpool.New(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect migration db: %v", err)
+	}
+	defer pool.Close()
+	for _, migrationPath := range migrationFiles {
+		migration, err := os.ReadFile(migrationPath)
+		if err != nil {
+			t.Fatalf("read migration %s: %v", filepath.Base(migrationPath), err)
+		}
+		if _, err := pool.Exec(ctx, string(migration)); err != nil {
+			t.Fatalf("apply migration %s: %v", filepath.Base(migrationPath), err)
+		}
+	}
+}
+
+func doJSON(t *testing.T, baseURL string, method string, path string, token string, payload any, expectedStatus int, out any) {
+	t.Helper()
+	var body io.Reader
+	if payload != nil {
+		raw, err := json.Marshal(payload)
+		if err != nil {
+			t.Fatalf("marshal payload: %v", err)
+		}
+		body = bytes.NewReader(raw)
+	}
+	req, err := http.NewRequest(method, baseURL+path, body)
+	if err != nil {
+		t.Fatalf("build request: %v", err)
+	}
+	if payload != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+	if token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("%s %s: %v", method, path, err)
+	}
+	defer resp.Body.Close()
+	raw, _ := io.ReadAll(resp.Body)
+	if resp.StatusCode != expectedStatus {
+		t.Fatalf("%s %s status=%d want=%d body=%s", method, path, resp.StatusCode, expectedStatus, string(raw))
+	}
+	if out != nil && len(raw) > 0 {
+		if err := json.Unmarshal(raw, out); err != nil {
+			t.Fatalf("decode %s %s response: %v body=%s", method, path, err, string(raw))
+		}
+	}
+}

apps/api/internal/httpapi/handlers.go

@@ -54,8 +54,12 @@ func (s *Server) register(w http.ResponseWriter, r *http.Request) {
 			writeError(w, http.StatusBadRequest, err.Error())
 			return
 		}
+		if errors.Is(err, store.ErrUserAlreadyExists) {
+			writeError(w, http.StatusConflict, err.Error())
+			return
+		}
 		s.logger.Error("register local user failed", "error", err)
-		writeError(w, http.StatusConflict, "user already exists or tenant is unavailable")
+		writeError(w, http.StatusInternalServerError, "register local user failed")
 		return
 	}
 	s.writeAuthResponse(w, http.StatusCreated, user)
@@ -230,6 +234,56 @@ func (s *Server) listUserGroups(w http.ResponseWriter, r *http.Request) {
 	writeJSON(w, http.StatusOK, map[string]any{"items": items})
 }
 
+func (s *Server) listAPIKeys(w http.ResponseWriter, r *http.Request) {
+	user, _ := auth.UserFromContext(r.Context())
+	items, err := s.store.ListAPIKeys(r.Context(), user)
+	if err != nil {
+		s.logger.Error("list api keys failed", "error", err)
+		writeError(w, http.StatusInternalServerError, "list api keys failed")
+		return
+	}
+	writeJSON(w, http.StatusOK, map[string]any{"items": items})
+}
+
+func (s *Server) createAPIKey(w http.ResponseWriter, r *http.Request) {
+	user, _ := auth.UserFromContext(r.Context())
+	var input store.CreateAPIKeyInput
+	if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
+		writeError(w, http.StatusBadRequest, "invalid json body")
+		return
+	}
+	created, err := s.store.CreateAPIKey(r.Context(), input, user)
+	if err != nil {
+		if errors.Is(err, store.ErrLocalUserRequired) {
+			writeError(w, http.StatusBadRequest, err.Error())
+			return
+		}
+		s.logger.Error("create api key failed", "error", err)
+		writeError(w, http.StatusInternalServerError, "create api key failed")
+		return
+	}
+	writeJSON(w, http.StatusCreated, created)
+}
+
+func (s *Server) disableAPIKey(w http.ResponseWriter, r *http.Request) {
+	user, _ := auth.UserFromContext(r.Context())
+	item, err := s.store.DisableAPIKey(r.Context(), r.PathValue("apiKeyID"), user)
+	if err == nil {
+		writeJSON(w, http.StatusOK, item)
+		return
+	}
+	if errors.Is(err, store.ErrLocalUserRequired) {
+		writeError(w, http.StatusBadRequest, err.Error())
+		return
+	}
+	if store.IsNotFound(err) {
+		writeError(w, http.StatusNotFound, "api key not found")
+		return
+	}
+	s.logger.Error("disable api key failed", "error", err)
+	writeError(w, http.StatusInternalServerError, "disable api key failed")
+}
+
 func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
 	var body map[string]any
 	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
@@ -275,6 +329,7 @@ func (s *Server) createTask(kind string) http.Handler {
 		task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
 			Kind:    kind,
 			Model:   model,
+			RunMode: runModeFromRequest(body),
 			Request: body,
 		}, user)
 		if err != nil {
@@ -322,24 +377,37 @@ func (s *Server) taskEvents(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("Connection", "keep-alive")
 
-	sendSSE(w, "task.accepted", map[string]any{
-		"taskId": task.ID,
-		"status": task.Status,
-	})
-	if flusher, ok := w.(http.Flusher); ok {
-		flusher.Flush()
-	}
-
-	timer := time.NewTimer(250 * time.Millisecond)
-	defer timer.Stop()
-
-	select {
-	case <-r.Context().Done():
+	events, err := s.store.ListTaskEvents(r.Context(), task.ID)
+	if err != nil {
+		s.logger.Error("list task events failed", "error", err)
 		return
-	case <-timer.C:
-		sendSSE(w, "task.placeholder", map[string]any{
-			"taskId":  task.ID,
-			"message": "runtime worker is not wired yet",
+	}
+	for _, event := range events {
+		sendSSE(w, event.EventType, event)
+		if flusher, ok := w.(http.Flusher); ok {
+			flusher.Flush()
+		}
+	}
+	if len(events) == 0 {
+		sendSSE(w, "task.accepted", map[string]any{
+			"taskId": task.ID,
+			"status": task.Status,
 		})
 	}
 }
+
+func runModeFromRequest(body map[string]any) string {
+	if value, ok := body["runMode"].(string); ok {
+		return value
+	}
+	if value, ok := body["mode"].(string); ok {
+		return value
+	}
+	if value, ok := body["simulation"].(bool); ok && value {
+		return "simulation"
+	}
+	if value, ok := body["testMode"].(bool); ok && value {
+		return "simulation"
+	}
+	return ""
+}

apps/api/internal/httpapi/server.go

@@ -24,6 +24,7 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
 		auth:   auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
 		logger: logger,
 	}
+	server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
 
 	mux := http.NewServeMux()
 	mux.HandleFunc("GET /healthz", server.health)
@@ -37,6 +38,9 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
 	mux.Handle("GET /api/v1/tenants", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listTenants)))
 	mux.Handle("GET /api/v1/users", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listUsers)))
 	mux.Handle("GET /api/v1/user-groups", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listUserGroups)))
+	mux.Handle("GET /api/v1/api-keys", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listAPIKeys)))
+	mux.Handle("POST /api/v1/api-keys", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.createAPIKey)))
+	mux.Handle("PATCH /api/v1/api-keys/{apiKeyID}/disable", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.disableAPIKey)))
 	mux.Handle("GET /api/v1/pricing/rules", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPricingRules)))
 	mux.Handle("POST /api/v1/pricing/estimate", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.estimatePricing)))
 	mux.Handle("GET /api/v1/platforms", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPlatforms)))
@@ -55,7 +59,7 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
 func (s *Server) cors(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		origin := r.Header.Get("Origin")
-		if origin != "" && (s.cfg.CORSAllowedOrigin == "*" || strings.EqualFold(origin, s.cfg.CORSAllowedOrigin)) {
+		if origin != "" && originAllowed(origin, s.cfg.CORSAllowedOrigin) {
 			w.Header().Set("Access-Control-Allow-Origin", origin)
 			w.Header().Set("Vary", "Origin")
 			w.Header().Set("Access-Control-Allow-Credentials", "true")
@@ -70,6 +74,16 @@ func (s *Server) cors(next http.Handler) http.Handler {
 	})
 }
 
+func originAllowed(origin string, allowed string) bool {
+	for _, item := range strings.Split(allowed, ",") {
+		item = strings.TrimSpace(item)
+		if item == "*" || strings.EqualFold(origin, item) {
+			return true
+		}
+	}
+	return false
+}
+
 func (s *Server) recover(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		defer func() {