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

scripts/graph_extraction_probe/budget.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"sync"
+
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+)
+
+func probeSendLimit(requested *int, environment string, envSet bool) (int, error) {
+	limit := requestbudget.GraphMaxSends
+	if requested != nil {
+		limit = *requested
+	}
+	if limit < 1 || limit > requestbudget.GraphMaxSends {
+		return 0, fmt.Errorf("invalid probe HTTP send cap")
+	}
+	if envSet && environment != strconv.Itoa(limit) {
+		return 0, fmt.Errorf("probe HTTP send cap differs between request and environment")
+	}
+	return limit, nil
+}
+
+// This tighter probe allowance delegates model, phase and deadline enforcement
+// to the existing graph guard; worker construction and defaults are untouched.
+type probeExecutionBudget struct {
+	*requestbudget.GraphExecution
+	mu      sync.Mutex
+	limit   int
+	stopped error
+	cancel  context.CancelCauseFunc
+}
+
+func bindProbeExecution(ctx context.Context, key, model string, limit int) (context.Context, *probeExecutionBudget, context.CancelFunc, error) {
+	if _, err := probeSendLimit(&limit, "", false); err != nil {
+		return ctx, nil, func() {}, err
+	}
+	ctx, base, baseCancel, err := requestbudget.BindGraphExecution(ctx, key, model)
+	if err != nil {
+		return ctx, nil, baseCancel, err
+	}
+	ctx, cancel := context.WithCancelCause(ctx)
+	guard := &probeExecutionBudget{GraphExecution: base, limit: limit, cancel: cancel}
+	return requestbudget.WithGuard(ctx, guard), guard, func() { cancel(nil); baseCancel() }, nil
+}
+
+func (g *probeExecutionBudget) Check(ctx context.Context) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if g.stopped != nil {
+		return g.stopped
+	}
+	return g.GraphExecution.Check(ctx)
+}
+
+func (g *probeExecutionBudget) Reserve(ctx context.Context, model, phase string) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if g.stopped != nil {
+		return g.stopped
+	}
+	if err := g.GraphExecution.Check(ctx); err != nil {
+		return err
+	}
+	if g.Reservations() >= g.limit {
+		g.stopped = fmt.Errorf("%w: probe HTTP send cap reached", requestbudget.ErrStopped)
+		g.cancel(g.stopped)
+		return g.stopped
+	}
+	return g.GraphExecution.Reserve(ctx, model, phase)
+}

scripts/graph_extraction_probe/budget_test.go

@@ -0,0 +1,196 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+)
+
+func TestProbeSendLimit(t *testing.T) {
+	one, two, three, zero, four := 1, 2, 3, 0, 4
+	for _, tc := range []struct {
+		name        string
+		requested   *int
+		environment string
+		envSet      bool
+		want        int
+	}{
+		{name: "default", want: 3},
+		{name: "minimum", requested: &one, want: 1},
+		{name: "two", requested: &two, environment: "2", envSet: true, want: 2},
+		{name: "explicit default", requested: &three, environment: "3", envSet: true, want: 3},
+		{name: "default environment", environment: "3", envSet: true, want: 3},
+		{name: "zero", requested: &zero},
+		{name: "excess", requested: &four},
+		{name: "environment mismatch", requested: &two, environment: "3", envSet: true},
+		{name: "environment without request", environment: "2", envSet: true},
+		{name: "empty environment", requested: &two, envSet: true},
+		{name: "padded environment", requested: &two, environment: "02", envSet: true},
+		{name: "decimal environment", requested: &two, environment: "2.0", envSet: true},
+		{name: "whitespace environment", requested: &two, environment: " 2", envSet: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := probeSendLimit(tc.requested, tc.environment, tc.envSet)
+			if got != tc.want || (err != nil) != (tc.want == 0) {
+				t.Fatalf("limit = %d, error = %v; want %d", got, err, tc.want)
+			}
+		})
+	}
+}
+
+type probeRoundTripper func(*http.Request) (*http.Response, error)
+
+func (f probeRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }
+
+func TestProbeTransportHonorsActualSendCap(t *testing.T) {
+	for _, limit := range []int{1, 2, requestbudget.GraphMaxSends} {
+		t.Run(fmt.Sprint(limit), func(t *testing.T) {
+			ctx, guard, cancel, err := bindProbeExecution(context.Background(), "offline", "model", limit)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer cancel()
+			ctx = requestbudget.WithModel(ctx, "model")
+			sends := 0
+			transport := requestbudget.Transport{Base: probeRoundTripper(func(req *http.Request) (*http.Response, error) {
+				sends++
+				if req.Context().Err() != nil {
+					t.Fatal("allowed response was canceled")
+				}
+				return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("usable response"))}, nil
+			})}
+			for i := 0; i < limit; i++ {
+				req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "https://offline.invalid", nil)
+				response, err := transport.RoundTrip(req)
+				if err != nil {
+					t.Fatalf("send %d: %v", i+1, err)
+				}
+				body, readErr := io.ReadAll(response.Body)
+				response.Body.Close()
+				if readErr != nil || string(body) != "usable response" || requestbudget.Check(ctx) != nil {
+					t.Fatal("last allowed response must remain usable")
+				}
+			}
+			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "https://offline.invalid", nil)
+			if _, err := transport.RoundTrip(req); !errors.Is(err, requestbudget.ErrStopped) {
+				t.Fatalf("excess send error = %v", err)
+			}
+			if sends != limit || guard.Reservations() != limit || !errors.Is(context.Cause(ctx), requestbudget.ErrStopped) {
+				t.Fatalf("sends = %d, reservations = %d, cause = %v", sends, guard.Reservations(), context.Cause(ctx))
+			}
+			if err := guard.Reserve(context.Background(), "model", requestbudget.GraphPhase); !errors.Is(err, requestbudget.ErrStopped) {
+				t.Fatal("stopped execution was revived")
+			}
+		})
+	}
+}
+
+func TestProbeBudgetPreservesScopeAndDeadline(t *testing.T) {
+	for _, tc := range []struct{ name, model, phase string }{
+		{"model", "other", requestbudget.GraphPhase},
+		{"phase", "model", "other"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, guard, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer cancel()
+			if err := guard.Reserve(ctx, tc.model, tc.phase); !errors.Is(err, requestbudget.ErrStopped) {
+				t.Fatalf("scope mismatch error = %v", err)
+			}
+			if guard.Reservations() != 0 || !errors.Is(context.Cause(ctx), requestbudget.ErrStopped) {
+				t.Fatal("scope mismatch consumed a send or did not stop execution")
+			}
+		})
+	}
+	t.Run("inherited deadline", func(t *testing.T) {
+		deadline := time.Now().Add(time.Minute)
+		parent, parentCancel := context.WithDeadline(context.Background(), deadline)
+		defer parentCancel()
+		ctx, guard, cancel, err := bindProbeExecution(parent, "offline", "model", 2)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cancel()
+		if got, ok := ctx.Deadline(); !ok || !got.Equal(deadline) {
+			t.Fatal("inherited deadline was reset")
+		}
+		if err := guard.Reserve(ctx, "model", requestbudget.GraphPhase); err != nil {
+			t.Fatal(err)
+		}
+		if got, _ := ctx.Deadline(); !got.Equal(deadline) {
+			t.Fatal("reservation reset the deadline")
+		}
+		parentCancel()
+		if err := guard.Reserve(ctx, "model", requestbudget.GraphPhase); !errors.Is(err, context.Canceled) {
+			t.Fatalf("cancellation error = %v", err)
+		}
+		if guard.Reservations() != 1 {
+			t.Fatal("canceled send consumed a reservation")
+		}
+	})
+	t.Run("default deadline", func(t *testing.T) {
+		started := time.Now()
+		ctx, _, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cancel()
+		deadline, ok := ctx.Deadline()
+		if !ok || deadline.Before(started.Add(requestbudget.GraphTimeout)) || deadline.After(time.Now().Add(requestbudget.GraphTimeout)) {
+			t.Fatal("probe did not retain the shared absolute timeout")
+		}
+	})
+	t.Run("already canceled", func(t *testing.T) {
+		parent, cancel := context.WithCancel(context.Background())
+		cancel()
+		if _, _, cleanup, err := bindProbeExecution(parent, "offline", "model", 2); err == nil {
+			cleanup()
+			t.Fatal("canceled execution was rebound")
+		}
+	})
+	t.Run("nested execution", func(t *testing.T) {
+		ctx, _, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cancel()
+		if _, _, cleanup, err := bindProbeExecution(ctx, "retry", "model", 2); err == nil {
+			cleanup()
+			t.Fatal("nested execution reset the allowance")
+		}
+	})
+}
+
+func TestProbeBudgetConcurrentReservationsStayBounded(t *testing.T) {
+	ctx, guard, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cancel()
+	var successful atomic.Int32
+	var group sync.WaitGroup
+	for i := 0; i < 16; i++ {
+		group.Add(1)
+		go func() {
+			defer group.Done()
+			if guard.Reserve(ctx, "model", requestbudget.GraphPhase) == nil {
+				successful.Add(1)
+			}
+		}()
+	}
+	group.Wait()
+	if successful.Load() != 2 || guard.Reservations() != 2 {
+		t.Fatalf("successful = %d; reservations = %d", successful.Load(), guard.Reservations())
+	}
+}

scripts/graph_extraction_probe/main.go

@@ -0,0 +1,170 @@
+// graph_extraction_probe calls the production extractor without writing graph data.
+// The model configuration and AES key arrive on stdin and are never reported.
+package main
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"time"
+	"unicode/utf8"
+
+	chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline"
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/utils"
+	"gopkg.in/yaml.v3"
+)
+
+type request struct {
+	ExecutionID   string              `json:"execution_id"`
+	MaxHTTPSends  *int                `json:"max_http_sends"`
+	AESKey        string              `json:"aes_key"`
+	Model         types.Model         `json:"model"`
+	Chunk         types.Chunk         `json:"chunk"`
+	ExtractConfig types.ExtractConfig `json:"extract_config"`
+}
+
+type chatDelegate interface{ chat.Chat }
+
+type probeCall struct {
+	Number           int              `json:"number"`
+	ElapsedMS        int64            `json:"elapsed_ms"`
+	MaxTokens        int              `json:"max_tokens"`
+	Thinking         *bool            `json:"thinking_requested"`
+	ResponseReceived bool             `json:"response_received"`
+	Failed           bool             `json:"failed"`
+	FinishReason     string           `json:"finish_reason"`
+	ContentRunes     int              `json:"content_runes"`
+	ReasoningRunes   int              `json:"reasoning_runes"`
+	Usage            types.TokenUsage `json:"usage"`
+}
+
+type recordingChat struct {
+	chatDelegate
+	response *types.ChatResponse
+	calls    int
+	records  []probeCall
+}
+
+func (c *recordingChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
+	started := time.Now()
+	response, err := c.chatDelegate.Chat(ctx, messages, options)
+	c.response = response
+	c.calls++
+	record := probeCall{
+		Number: c.calls, ElapsedMS: time.Since(started).Milliseconds(),
+		ResponseReceived: response != nil, Failed: err != nil,
+	}
+	if options != nil {
+		record.MaxTokens = options.MaxTokens
+		if options.Thinking != nil {
+			thinking := *options.Thinking
+			record.Thinking = &thinking
+		}
+	}
+	if response != nil {
+		record.FinishReason = response.FinishReason
+		record.ContentRunes = utf8.RuneCountInString(response.Content)
+		record.ReasoningRunes = utf8.RuneCountInString(response.ReasoningContent)
+		record.Usage = response.Usage
+	}
+	c.records = append(c.records, record)
+	return response, err
+}
+
+func run() error {
+	var input request
+	if err := json.NewDecoder(os.Stdin).Decode(&input); err != nil {
+		return fmt.Errorf("invalid probe input")
+	}
+	capValue, capSet := os.LookupEnv("GRAPH_PROBE_MAX_SENDS")
+	maxSends, err := probeSendLimit(input.MaxHTTPSends, capValue, capSet)
+	if err != nil {
+		return err
+	}
+	if input.ExecutionID == "" || input.Chunk.ID == "" || input.Chunk.Content == "" || input.Model.ID == "" || input.Model.Name == "" {
+		return fmt.Errorf("probe requires a source chunk and model")
+	}
+	if err := os.Setenv("SYSTEM_AES_KEY", input.AESKey); err != nil {
+		return err
+	}
+	input.AESKey = ""
+	key, err := utils.DecryptStoredSecret(input.Model.Parameters.APIKey)
+	if err != nil {
+		return fmt.Errorf("cannot decrypt configured model credential")
+	}
+	input.Model.Parameters.APIKey = key
+	defer os.Unsetenv("SYSTEM_AES_KEY")
+	chatConfig := chat.ConfigFromModel(&input.Model, "", "")
+	model, err := chat.NewChat(chatConfig, nil)
+	if err != nil {
+		return fmt.Errorf("cannot initialize configured remote model")
+	}
+	base, err := os.ReadFile("config/config.yaml")
+	if err != nil {
+		return err
+	}
+	var config struct {
+		Extract struct {
+			Graph struct {
+				Description string `yaml:"description"`
+			} `yaml:"extract_graph"`
+		} `yaml:"extract"`
+	}
+	if err := yaml.Unmarshal(base, &config); err != nil {
+		return err
+	}
+	recorder := &recordingChat{chatDelegate: model}
+	extractor := chatpipeline.NewGraphExtractor(recorder, &types.PromptTemplateStructured{
+		Description: types.AppendCustomPromptInstructions(config.Extract.Graph.Description, input.ExtractConfig.CustomInstructions, "graph_extraction"),
+		Tags:        input.ExtractConfig.Tags,
+	})
+	ctx, guard, cancel, err := bindProbeExecution(context.Background(), input.ExecutionID, input.Model.ID, maxSends)
+	if err != nil {
+		return err
+	}
+	defer cancel()
+	started := time.Now()
+	graph, extractionErr := extractor.Extract(ctx, input.Chunk.Content)
+	hash := sha256.Sum256([]byte(input.Chunk.Content))
+	report := map[string]any{
+		"checked_at": started.UTC().Format(time.RFC3339), "elapsed_ms": time.Since(started).Milliseconds(),
+		"chunk_id": input.Chunk.ID, "source_sha256": hex.EncodeToString(hash[:]),
+		"model_name": input.Model.Name, "allowed_relation_types": input.ExtractConfig.Tags,
+		"passed": extractionErr == nil, "graph": graph, "production_write": false,
+		"model_calls":                recorder.calls,
+		"execution_id":               input.ExecutionID,
+		"http_send_reservations":     guard.Reservations(),
+		"http_send_limit":            maxSends,
+		"call_diagnostics":           recorder.records,
+		"effective_thinking_control": chat.EffectiveThinkingControl(chatConfig),
+		"diagnostics_version":        2,
+	}
+	if recorder.response != nil {
+		report["finish_reason"] = recorder.response.FinishReason
+		report["usage"] = recorder.response.Usage
+		if extractionErr != nil {
+			report["rejected_model_content"] = recorder.response.Content
+		}
+	}
+	if extractionErr != nil {
+		report["error"] = extractionErr.Error()
+	}
+	encoded, err := json.Marshal(report)
+	if err != nil {
+		return err
+	}
+	fmt.Printf("GRAPH_EXTRACTION_PROBE_RESULT=%s\n", encoded)
+	return nil
+}
+
+func main() {
+	if err := run(); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+}

scripts/graph_extraction_probe/main_test.go

@@ -0,0 +1,77 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"strings"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+type fakeProbeChat struct {
+	chatDelegate
+	response *types.ChatResponse
+	err      error
+}
+
+func (f *fakeProbeChat) Chat(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error) {
+	return f.response, f.err
+}
+
+func TestRecordingChatPreservesPerCallDiagnosticsWithoutContent(t *testing.T) {
+	fake := &fakeProbeChat{response: &types.ChatResponse{
+		Content: "answer", ReasoningContent: "private reasoning", FinishReason: "length",
+		Usage: types.TokenUsage{PromptTokens: 10, CompletionTokens: 4, TotalTokens: 14},
+	}}
+	recorder := &recordingChat{chatDelegate: fake}
+	thinking := false
+	options := &chat.ChatOptions{MaxTokens: 4096, Thinking: &thinking}
+	first, err := recorder.Chat(context.Background(), nil, options)
+	if err != nil || first != fake.response {
+		t.Fatal("recorder changed the response")
+	}
+	thinking = true
+	fake.response = &types.ChatResponse{Content: "\u4e2d\u6587", FinishReason: "stop"}
+	if _, err := recorder.Chat(context.Background(), nil, options); err != nil {
+		t.Fatal(err)
+	}
+	if recorder.calls != 2 || len(recorder.records) != 2 {
+		t.Fatal("missing call diagnostics")
+	}
+	a, b := recorder.records[0], recorder.records[1]
+	if a.Number != 1 || a.MaxTokens != 4096 || a.Thinking == nil || *a.Thinking ||
+		!a.ResponseReceived || a.Failed || a.FinishReason != "length" ||
+		a.ContentRunes != 6 || a.ReasoningRunes != 17 || a.Usage.TotalTokens != 14 {
+		t.Fatalf("unexpected first call metadata: %+v", a)
+	}
+	if b.Number != 2 || b.Thinking == nil || !*b.Thinking || b.ContentRunes != 2 || b.ReasoningRunes != 0 {
+		t.Fatalf("unexpected second call metadata: %+v", b)
+	}
+	encoded, err := json.Marshal(recorder.records)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if strings.Contains(string(encoded), "private reasoning") || strings.Contains(string(encoded), "answer") {
+		t.Fatal("diagnostics exposed response content")
+	}
+}
+
+func TestRecordingChatRecordsNoResponseWithoutLeakingError(t *testing.T) {
+	want := errors.New("private upstream error")
+	recorder := &recordingChat{chatDelegate: &fakeProbeChat{err: want}}
+	response, err := recorder.Chat(context.Background(), nil, nil)
+	if response != nil || err != want || len(recorder.records) != 1 {
+		t.Fatal("recorder changed the failure")
+	}
+	record := recorder.records[0]
+	if record.ResponseReceived || !record.Failed || record.Thinking != nil || record.ContentRunes != 0 || record.ReasoningRunes != 0 {
+		t.Fatalf("unexpected failure metadata: %+v", record)
+	}
+	encoded, _ := json.Marshal(record)
+	if strings.Contains(string(encoded), want.Error()) {
+		t.Fatal("diagnostics exposed upstream error")
+	}
+}