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

internal/workbench/instructions_test.go

@@ -0,0 +1,69 @@
+package workbench
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+)
+
+type instructionStub struct {
+	text string
+	err  error
+}
+
+type disabledInstructionAgentStub struct{ interfaces.CustomAgentService }
+
+func (disabledInstructionAgentStub) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
+	tenant, _ := types.TenantIDFromContext(ctx)
+	return &types.CustomAgent{ID: id, TenantID: tenant, Visibility: "global", Enabled: false, Config: types.CustomAgentConfig{ModelID: "model"}}, nil
+}
+
+func TestDisabledAgentCannotGenerateWithoutInstructionResolver(t *testing.T) {
+	s, ctx, _, model := generationHarness(t)
+	s.agents = disabledInstructionAgentStub{}
+	if _, err := s.complete(ctx, "disabled", "system", "prompt"); !errors.Is(err, ErrInvalid) {
+		t.Fatalf("expected disabled Agent rejection, got %v", err)
+	}
+	if len(model.messages) != 0 {
+		t.Fatal("disabled Agent reached provider")
+	}
+}
+
+func (s instructionStub) ResolveAgentInstructions(context.Context, *types.CustomAgent) (string, error) {
+	return s.text, s.err
+}
+
+func TestMountedInstructionsParticipateInGenerationAndBudget(t *testing.T) {
+	s, ctx, _, model := generationHarness(t)
+	s.SetInstructionResolver(instructionStub{text: "技能规则：先列判断依据，再给案例。"})
+	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(model.messages[0].Content, "技能规则：先列判断依据") {
+		t.Fatal("mounted instruction missing from provider request")
+	}
+	s.models = generationModelStub{model: model, parameters: types.ModelParameters{ContextWindow: 4096, MaxOutputTokens: 512}}
+	s.SetInstructionResolver(instructionStub{text: strings.Repeat("很长的技能规则", 3000)})
+	model.messages = nil
+	if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrInvalid) {
+		t.Fatalf("expected budget rejection, got %v", err)
+	}
+	if len(model.messages) != 0 {
+		t.Fatal("oversized skills reached model")
+	}
+}
+
+func TestUnavailableInstructionStopsGeneration(t *testing.T) {
+	s, ctx, _, model := generationHarness(t)
+	s.SetInstructionResolver(instructionStub{err: ErrForbidden})
+	if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrForbidden) {
+		t.Fatalf("got %v", err)
+	}
+	if len(model.messages) != 0 {
+		t.Fatal("generation silently ignored unavailable instruction")
+	}
+}