ai-gateway-go/internal/workbench/workbench_integration_test.go

package workbench

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"aigateway.local/core/internal/apikey"
	"aigateway.local/core/internal/gateway"
	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/cryptox"
	"aigateway.local/core/internal/platform/database"
	tracepkg "aigateway.local/core/internal/trace"
)

func TestWorkbenchPostgreSQLLifecycle(t *testing.T) {
	databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL")
	if databaseURL == "" {
		t.Skip("WORKBENCH_TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8, MinConns: 0})
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	actorID := "11111111-1111-4111-8111-111111111111"
	_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role) VALUES($1,'m4-test','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID)
	if err != nil {
		t.Fatal(err)
	}
	cleanup := func() {
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE request_id='m4-runtime'; DELETE FROM gateway.notification_channels WHERE name='m4-webhook'; DELETE FROM gateway.applications WHERE code='m4_app'; DELETE FROM gateway.tool_definitions WHERE code='m4_lookup'; DELETE FROM gateway.knowledge_bases WHERE name='m4-integration-kb'; DELETE FROM gateway.prompt_templates WHERE name='m4-integration-prompt'`)
	}
	cleanup()
	defer cleanup()
	assets := NewService(pool)
	prompt, err := assets.CreatePrompt(ctx, PromptInput{Name: "m4-integration-prompt", Description: "test", Enabled: true, Content: "请回答 {{question}}", Variables: []Variable{{Name: "question", Required: true}}}, actorID)
	if err != nil {
		t.Fatal(err)
	}
	if prompt.Current == nil || prompt.Current.Version != 1 {
		t.Fatalf("unexpected prompt: %#v", prompt)
	}
	version, err := assets.AddPromptVersion(ctx, prompt.ID, "新版 {{question}}", []Variable{{Name: "question", Required: true}}, "v2", actorID, true)
	if err != nil || version.Version != 2 {
		t.Fatalf("version=%#v err=%v", version, err)
	}
	_, rendered, err := assets.RenderPromptByName(ctx, prompt.Name, map[string]any{"question": "可扩展吗"})
	if err != nil || rendered != "新版 可扩展吗" {
		t.Fatalf("render=%q err=%v", rendered, err)
	}
	kb, err := assets.SaveKnowledgeBase(ctx, KnowledgeBase{Name: "m4-integration-kb", Description: "test", RetrievalMode: "postgres_fts", ChunkSize: 200, ChunkOverlap: 20, Enabled: true}, actorID, true)
	if err != nil {
		t.Fatal(err)
	}
	doc, err := assets.AddKnowledgeDocument(ctx, kb.ID, "Go 架构说明", "text", "", "Go 网关采用不可变运行时快照。\n\nPostgreSQL 是权威配置存储。", actorID)
	if err != nil {
		t.Fatal(err)
	}
	if doc.ChunkCount == 0 {
		t.Fatal("expected chunks")
	}
	hits, err := NewRetriever(assets, nil).Search(ctx, kb.ID, "不可变运行时快照", 4)
	if err != nil || len(hits) == 0 {
		t.Fatalf("hits=%#v err=%v", hits, err)
	}
	key := base64.StdEncoding.EncodeToString(make([]byte, 32))
	cipher, err := cryptox.NewKeyring(key, 1, "", "m4-test-tools")
	if err != nil {
		t.Fatal(err)
	}
	receivedAuth := ""
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		receivedAuth = r.Header.Get("Authorization")
		_ = json.NewEncoder(w).Encode(map[string]any{"ok": true})
	}))
	defer upstream.Close()
	tools := NewToolService(assets, cipher, true)
	tool, err := tools.Save(ctx, "", ToolInput{Code: "m4_lookup", Name: "查询", EndpointURL: upstream.URL, HTTPMethod: "POST", Headers: map[string]string{"Authorization": "Bearer secret"}, InputSchema: json.RawMessage(`{"type":"object","required":["q"],"properties":{"q":{"type":"string"}}}`), TimeoutSeconds: 5, Enabled: true}, actorID, true)
	if err != nil {
		t.Fatal(err)
	}
	result, err := tools.Execute(ctx, tool, map[string]any{"q": "gateway"}, "", "m4-test")
	if err != nil {
		t.Fatal(err)
	}
	if result["status_code"] != http.StatusOK || receivedAuth != "Bearer secret" {
		t.Fatalf("result=%#v auth=%q", result, receivedAuth)
	}
	app, err := assets.SaveApplication(ctx, Application{Code: "m4_app", Name: "M4 App", Status: "draft", DraftConfig: ApplicationConfig{Model: "test-model", PromptTemplateID: prompt.ID, KnowledgeBaseIDs: []string{kb.ID}, ToolIDs: []string{tool.ID}, RetrievalTopK: 4, Temperature: .2, MaxToolRounds: 2}}, actorID, true)
	if err != nil {
		t.Fatal(err)
	}
	published, err := assets.PublishApplication(ctx, app.ID, "first", actorID)
	if err != nil {
		t.Fatal(err)
	}
	if published.PublishedVersion == nil || *published.PublishedVersion != 1 || published.Status != "active" {
		t.Fatalf("unexpected published app: %#v", published)
	}
	governed := false
	fakeGateway := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var payload map[string]any
		if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
			t.Error(err)
		}
		messages, _ := payload["messages"].([]any)
		if len(messages) > 0 && strings.Contains(toString(messages[0]), "不可变运行时快照") {
			governed = true
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "完成"}}}})
	})
	runtime := NewRuntimeHTTPHandler(assets, tools, NewRetriever(assets, nil), staticPrincipalAuthenticator{}, fakeGateway, MarketplaceDeps{})
	traceStore := tracepkg.NewStore(pool)
	runtime.SetTraceStore(traceStore)
	runtimeRequest := httptest.NewRequest(http.MethodPost, "/v1/applications/m4_app/chat/completions", bytes.NewBufferString(`{"messages":[{"role":"user","content":"不可变运行时快照是什么？"}],"variables":{"question":"架构"}}`))
	runtimeRequest.Header.Set("Authorization", "Bearer test")
	runtimeRequest = runtimeRequest.WithContext(gateway.WithRequestID(runtimeRequest.Context(), "m4-runtime"))
	runtimeResponse := httptest.NewRecorder()
	runtime.ServeHTTP(runtimeResponse, runtimeRequest)
	if runtimeResponse.Code != http.StatusOK || !governed || !strings.Contains(runtimeResponse.Body.String(), `"application"`) {
		t.Fatalf("runtime status=%d governed=%v body=%s", runtimeResponse.Code, governed, runtimeResponse.Body.String())
	}
	traces, err := traceStore.List(ctx, tracepkg.Filter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), RequestID: "m4-runtime", Limit: 10})
	if err != nil || len(traces) != 1 || traces[0].TraceType != "application" || traces[0].ModelCallCount != 1 {
		t.Fatalf("runtime trace=%+v err=%v", traces, err)
	}
	detail, err := traceStore.Get(ctx, traces[0].ID)
	if err != nil || len(detail.Spans) < 2 {
		t.Fatalf("runtime trace detail=%+v err=%v", detail, err)
	}

	signed := ""
	webhook := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		signed = r.Header.Get("X-Gateway-Signature")
		w.WriteHeader(http.StatusNoContent)
	}))
	defer webhook.Close()
	notifications := NewNotificationService(assets, cipher, true)
	secret := "m4-signing-secret"
	channel, err := notifications.SaveChannel(ctx, "", NotificationInput{Name: "m4-webhook", WebhookURL: webhook.URL, SigningSecret: &secret, EventPatterns: []string{"application.*"}, Enabled: true}, actorID, true)
	if err != nil {
		t.Fatal(err)
	}
	eventID, _ := newUUID()
	delivery, err := notifications.ensureDelivery(ctx, channel, eventID, "application.published", json.RawMessage(`{"version":1}`))
	if err != nil {
		t.Fatal(err)
	}
	if err = notifications.deliver(ctx, channel, delivery); err != nil || !strings.HasPrefix(signed, "sha256=") {
		t.Fatalf("delivery err=%v signature=%q", err, signed)
	}
}

type staticPrincipalAuthenticator struct{}

func (staticPrincipalAuthenticator) AuthenticatePrincipal(context.Context, string) (apikey.Principal, error) {
	return apikey.Principal{}, nil
}