xinmi/internal/workbench/agent_snapshot_test.go

package workbench

import (
	"context"
	"errors"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

type snapshotInstructions struct {
	text string
	err  error
}

func (s *snapshotInstructions) ResolveAgentInstructions(context.Context, *types.CustomAgent) (string, error) {
	return s.text, s.err
}

func TestFrozenAgentRetainsInstructionsButRechecksSkillAccess(t *testing.T) {
	s, ctx, _, model := generationHarness(t)
	s.agents = generationAgentStub{config: types.CustomAgentConfig{SystemPrompt: "原人格", InstructionSkillIDs: []string{"skill"}, Temperature: 0.3, MaxCompletionTokens: 1000}}
	instructions := &snapshotInstructions{text: "原技能"}
	s.SetInstructionResolver(instructions)
	snapshots := map[string]AgentSnapshot{}
	ctx = context.WithValue(ctx, agentSnapshotsKey{}, snapshots)
	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
		t.Fatal(err)
	}
	s.agents = generationAgentStub{config: types.CustomAgentConfig{SystemPrompt: "新人格", Temperature: 0.8, MaxCompletionTokens: 4000}}
	instructions.text = "新技能"
	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(model.messages[0].Content, "原人格") || !strings.Contains(model.messages[0].Content, "原技能") || strings.Contains(model.messages[0].Content, "新技能") || model.options.Temperature != 0.3 || model.options.MaxCompletionTokens != 1000 {
		t.Fatal("a saved task did not preserve its Agent and skill inputs")
	}
	instructions.err = ErrForbidden
	if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrForbidden) {
		t.Fatalf("revoked skill must not be usable via snapshot: %v", err)
	}
}

type snapshotMemory struct {
	interfaces.MemoryService
	user  string
	calls int
}

func (m *snapshotMemory) Recall(ctx context.Context, _ string) interfaces.MemoryRecall {
	m.user, _ = types.UserIDFromContext(ctx)
	m.calls++
	return interfaces.MemoryRecall{Prompt: "偏好：简洁表达"}
}

func TestTaskMemoryUsesCurrentPrincipalAndAgentOptOut(t *testing.T) {
	s, ctx, _, model := generationHarness(t)
	m := &snapshotMemory{}
	s.SetMemoryService(m)
	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
		t.Fatal(err)
	}
	if m.user != "writer" || !strings.Contains(model.messages[0].Content, "简洁表达") {
		t.Fatal("missing scoped task memory")
	}
	disabled := false
	s.agents = generationAgentStub{config: types.CustomAgentConfig{MemoryEnabled: &disabled}}
	if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
		t.Fatal(err)
	}
	if m.calls != 1 || strings.Contains(model.messages[0].Content, "简洁表达") {
		t.Fatal("agent opt-out ignored")
	}
}