WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,253 @@
+package chat
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+)
+
+type chatBudgetGuard struct {
+	mu             sync.Mutex
+	limit          int
+	models, phases []string
+	checkErr       error
+}
+
+func (g *chatBudgetGuard) Key() string { return "local-test" }
+func (g *chatBudgetGuard) Check(ctx context.Context) error {
+	if g.checkErr != nil {
+		return g.checkErr
+	}
+	return ctx.Err()
+}
+func (g *chatBudgetGuard) Reserve(ctx context.Context, model, phase string) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	if len(g.models) >= g.limit {
+		return requestbudget.ErrStopped
+	}
+	g.models, g.phases = append(g.models, model), append(g.phases, phase)
+	return nil
+}
+
+func chatBudgetContext(g *chatBudgetGuard) context.Context {
+	return requestbudget.WithModel(requestbudget.WithPhase(requestbudget.WithGuard(context.Background(), g), "wiki_generate"), "caller-forged-model")
+}
+
+func TestRequestBudgetChatSDKRawAndStreaming(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	secutils.ResetSSRFWhitelistForTest()
+	t.Cleanup(secutils.ResetSSRFWhitelistForTest)
+	for _, raw := range []bool{false, true} {
+		for _, stream := range []bool{false, true} {
+			for _, limit := range []int{0, 1} {
+				t.Run(fmt.Sprintf("raw=%t/stream=%t/limit=%d", raw, stream, limit), func(t *testing.T) {
+					var sends atomic.Int32
+					server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+						sends.Add(1)
+						if stream {
+							w.Header().Set("Content-Type", "text/event-stream")
+							fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n")
+						} else {
+							w.Header().Set("Content-Type", "application/json")
+							fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`)
+						}
+					}))
+					defer server.Close()
+					c, err := NewRemoteAPIChat(&ChatConfig{BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"})
+					if err != nil {
+						t.Fatal(err)
+					}
+					opts := &ChatOptions{CacheRetention: CacheRetentionNone}
+					if raw {
+						opts.CacheRetention, opts.PromptCacheKey = CacheRetentionShort, "force-raw"
+					}
+					g := &chatBudgetGuard{limit: limit}
+					ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
+					defer cancel()
+					messages := []Message{{Role: "user", Content: "local test"}}
+					_, _, actualRaw, err := c.buildOutbound(ctx, messages, opts, stream)
+					if err != nil || actualRaw != raw {
+						t.Fatalf("test route mismatch: %t %v", actualRaw, err)
+					}
+					if stream {
+						var responses <-chan types.StreamResponse
+						responses, err = c.ChatStream(ctx, messages, opts)
+						if err == nil {
+							for range responses {
+							}
+						}
+					} else {
+						_, err = c.Chat(ctx, messages, opts)
+					}
+					if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
+						t.Fatalf("expected stop: %v", err)
+					}
+					if limit == 1 && err != nil {
+						t.Fatal(err)
+					}
+					if int(sends.Load()) != limit || len(g.models) != limit {
+						t.Fatalf("unaccounted sends: %d models=%v", sends.Load(), g.models)
+					}
+					if limit == 1 && (g.models[0] != "concrete-chat" || g.phases[0] != "wiki_generate") {
+						t.Fatalf("identity/phase lost: %v %v", g.models, g.phases)
+					}
+				})
+			}
+		}
+	}
+}
+
+func TestRequestBudgetMultimodalFallbackCountsEachSend(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	secutils.ResetSSRFWhitelistForTest()
+	t.Cleanup(secutils.ResetSSRFWhitelistForTest)
+	for _, stream := range []bool{false, true} {
+		for _, limit := range []int{1, 2} {
+			t.Run(fmt.Sprintf("stream=%t/limit=%d", stream, limit), func(t *testing.T) {
+				var sends atomic.Int32
+				server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					if sends.Add(1) == 1 {
+						w.Header().Set("Content-Type", "application/json")
+						w.WriteHeader(http.StatusBadRequest)
+						fmt.Fprint(w, `{"error":{"message":"image input not supported","type":"invalid_request_error"}}`)
+						return
+					}
+					if stream {
+						w.Header().Set("Content-Type", "text/event-stream")
+						fmt.Fprint(w, "data: [DONE]\n\n")
+					} else {
+						w.Header().Set("Content-Type", "application/json")
+						fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"}}]}`)
+					}
+				}))
+				defer server.Close()
+				c, err := NewRemoteAPIChat(&ChatConfig{BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"})
+				if err != nil {
+					t.Fatal(err)
+				}
+				g := &chatBudgetGuard{limit: limit}
+				ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
+				defer cancel()
+				messages := []Message{{Role: "user", Content: "local test", Images: []string{"data:image/png;base64,AA=="}}}
+				opts := &ChatOptions{CacheRetention: CacheRetentionNone}
+				if stream {
+					var responses <-chan types.StreamResponse
+					responses, err = c.ChatStream(ctx, messages, opts)
+					if err == nil {
+						for range responses {
+						}
+					}
+				} else {
+					_, err = c.Chat(ctx, messages, opts)
+				}
+				if limit == 1 && !errors.Is(err, requestbudget.ErrStopped) {
+					t.Fatalf("fallback escaped budget: %v", err)
+				}
+				if limit == 2 && err != nil {
+					t.Fatal(err)
+				}
+				if int(sends.Load()) != limit || len(g.models) != limit {
+					t.Fatalf("fallback sends=%d reservations=%d", sends.Load(), len(g.models))
+				}
+				for i := range g.models {
+					if g.models[i] != "concrete-chat" || g.phases[i] != "wiki_generate" {
+						t.Fatal("fallback context lost")
+					}
+				}
+			})
+		}
+	}
+}
+
+func TestRequestBudgetChatFactoryFailsClosed(t *testing.T) {
+	g := &chatBudgetGuard{limit: 1}
+	ctx := chatBudgetContext(g)
+	c, err := NewChat(&ChatConfig{Source: types.ModelSourceLocal, ModelID: "local-model"}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := c.Chat(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("local adapter invoked: %v", err)
+	}
+	if _, err := c.ChatStream(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("local stream invoked: %v", err)
+	}
+	if len(g.models) != 0 {
+		t.Fatal("unsupported adapter reserved a send")
+	}
+	unknown, err := wrapChatRequestBudget(&fakeChat{id: "unknown"}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := unknown.Chat(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("unknown adapter allowed: %v", err)
+	}
+}
+
+func TestRequestBudgetAnthropicFactory(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	secutils.ResetSSRFWhitelistForTest()
+	t.Cleanup(secutils.ResetSSRFWhitelistForTest)
+	for _, stream := range []bool{false, true} {
+		for _, limit := range []int{0, 1} {
+			t.Run(fmt.Sprintf("stream=%t/limit=%d", stream, limit), func(t *testing.T) {
+				var sends atomic.Int32
+				server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					sends.Add(1)
+					if stream {
+						w.Header().Set("Content-Type", "text/event-stream")
+						fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
+					} else {
+						w.Header().Set("Content-Type", "application/json")
+						fmt.Fprint(w, `{"type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn"}`)
+					}
+				}))
+				defer server.Close()
+				c, err := NewChat(&ChatConfig{Source: types.ModelSourceRemote, BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "claude-test", APIKey: "local-only", Provider: "anthropic"}, nil)
+				if err != nil {
+					t.Fatal(err)
+				}
+				g := &chatBudgetGuard{limit: limit}
+				ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
+				defer cancel()
+				messages := []Message{{Role: "user", Content: "local test"}}
+				if stream {
+					var responses <-chan types.StreamResponse
+					responses, err = c.ChatStream(ctx, messages, nil)
+					if err == nil {
+						for range responses {
+						}
+					}
+				} else {
+					_, err = c.Chat(ctx, messages, nil)
+				}
+				if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
+					t.Fatalf("guard not applied: %v", err)
+				}
+				if limit == 1 && err != nil {
+					t.Fatal(err)
+				}
+				if int(sends.Load()) != limit || len(g.models) != limit {
+					t.Fatalf("sends=%d models=%v", sends.Load(), g.models)
+				}
+				if limit == 1 && (g.models[0] != "concrete-chat" || g.phases[0] != "wiki_generate") {
+					t.Fatalf("context lost: %v %v", g.models, g.phases)
+				}
+			})
+		}
+	}
+}