xinmi/internal/application/service/digital_human_test.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
)

func TestDigitalHumanIdentityScopeDoesNotImpersonateResourceOwners(t *testing.T) {
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
	ctx = context.WithValue(ctx, types.UserIDContextKey, "owner")
	read, err := digitalHumanReadScope(ctx)
	require.NoError(t, err)
	require.False(t, read.ReadAll)
	ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleAdmin)
	read, err = digitalHumanReadScope(ctx)
	require.NoError(t, err)
	require.True(t, read.ReadAll)
	write, err := digitalHumanScope(ctx)
	require.NoError(t, err)
	require.False(t, write.ReadAll)
	require.Equal(t, uint64(7), write.TenantID)
	require.Equal(t, "owner", write.OwnerID)
	_, err = digitalHumanScope(context.WithValue(ctx, types.UserIDContextKey, "system-7"))
	require.ErrorIs(t, err, types.ErrDigitalHumanIdentity)
}

func TestDigitalHumanAdministratorManagesOtherOwnersWithinTenant(t *testing.T) {
	svc, owner, human, _, _ := setupDigitalHumanWeChat(t)
	admin := context.WithValue(owner, types.UserIDContextKey, "admin")
	admin = context.WithValue(admin, types.TenantRoleContextKey, types.TenantRoleAdmin)
	updated, err := svc.Update(admin, human.ID, DigitalHumanInput{Name: "Reviewed teacher", Revision: human.Revision, Persona: map[string]string{"tone": "new tone"}})
	require.NoError(t, err)
	require.Equal(t, "owner", updated.OwnerID)
	require.Equal(t, "owner", updated.Versions[0].OwnerID)
	reviewed, err := svc.Review(admin, human.ID, updated.LatestVersion, true, "approved")
	require.NoError(t, err)
	require.Equal(t, "active", reviewed.Status)
	shared := accesspolicy.WithSnapshot(context.WithValue(owner, types.UserIDContextKey, "recipient"), &accesspolicy.Snapshot{Publications: map[string]bool{"digital-humans:" + human.ID: true}})
	_, err = svc.Get(shared, human.ID)
	require.NoError(t, err)
	_, err = svc.Update(shared, human.ID, DigitalHumanInput{Name: "not allowed", Revision: reviewed.Revision})
	require.Error(t, err)
	require.Error(t, svc.Delete(shared, human.ID))
	foreignAdmin := context.WithValue(admin, types.TenantIDContextKey, uint64(8))
	require.Error(t, svc.Delete(foreignAdmin, human.ID))
	require.NoError(t, svc.Delete(admin, human.ID))
	_, err = svc.Get(owner, human.ID)
	require.ErrorIs(t, err, types.ErrDigitalHumanNotFound)
}

func TestDigitalHumanSourcesMergeAndFullUnicodeCoverage(t *testing.T) {
	previous := []types.DigitalHumanSource{{Name: "original", Text: strings.Repeat("文", 12003)}}
	added := []types.DigitalHumanSource{{Name: "duplicate", Text: "  " + previous[0].Text + "\n"}, {Name: "new", Text: "next source"}}
	sources, err := mergeDigitalHumanSources(previous, added)
	require.NoError(t, err)
	require.Len(t, sources, 2)
	batches := digitalHumanBatches(sources)
	require.Len(t, batches, 2)
	require.Equal(t, previous[0].Text, batches[0][0].Text+batches[1][0].Text)
	require.Equal(t, "next source", batches[1][1].Text)
	for _, batch := range batches {
		require.True(t, utf8.ValidString(batch[0].Text))
		require.LessOrEqual(t, utf8.RuneCountInString(batch[0].Text), 12000)
	}
	_, err = mergeDigitalHumanSources(nil, []types.DigitalHumanSource{{Text: strings.Repeat("x", digitalHumanSourceLimit+1)}})
	require.Error(t, err)
	var inputErr *types.DigitalHumanInputError
	require.ErrorAs(t, err, &inputErr)
}

func TestDigitalHuman99FilesAndIncrementalArchive(t *testing.T) {
	var sources []types.DigitalHumanSource
	var ids []string
	for i := 0; i < 99; i++ {
		ids = append(ids, fmt.Sprintf("upload-%d", i))
		sources = append(sources, types.DigitalHumanSource{Name: fmt.Sprintf("file-%d", i), Text: strings.Repeat(fmt.Sprintf("文%d ", i), 4000)})
	}
	require.NoError(t, validateDigitalHumanDistillCount(DigitalHumanDistillInput{UploadIDs: ids}))
	require.NoError(t, validateDigitalHumanDistillCount(DigitalHumanDistillInput{UploadIDs: ids, Sources: []types.DigitalHumanSource{{Text: "extra text"}}}))
	require.Error(t, validateDigitalHumanDistillCount(DigitalHumanDistillInput{UploadIDs: append(ids, "one-too-many")}))
	require.Error(t, validateDigitalHumanDistillCount(DigitalHumanDistillInput{UploadIDs: ids, Files: []DigitalHumanUpload{{Name: "extra.txt"}}}))
	require.Error(t, validateDigitalHumanDistillCount(DigitalHumanDistillInput{UploadIDs: ids, Sources: []types.DigitalHumanSource{{Text: "first"}, {Text: "second"}}}))
	merged, err := mergeDigitalHumanSources(sources, []types.DigitalHumanSource{{Name: "new", Text: "new evidence"}, {Name: "later", Text: "later evidence"}})
	require.NoError(t, err)
	require.Len(t, merged, 101)
	require.Len(t, newDigitalHumanSources(sources, merged), 2)
	reconstructed := make(map[string]string)
	for _, batch := range digitalHumanBatches(merged) {
		count := 0
		for _, source := range batch {
			require.True(t, utf8.ValidString(source.Text))
			count += utf8.RuneCountInString(source.Text)
			reconstructed[source.Name] += source.Text
		}
		require.LessOrEqual(t, count, digitalHumanBatchRunes)
	}
	for _, source := range merged {
		require.Equal(t, source.Text, reconstructed[source.Name], source.Name)
	}
	small := []types.DigitalHumanSource{{Name: "one", Text: "one"}, {Name: "two", Text: "two"}}
	require.Len(t, digitalHumanBatches(small), 1, "small files should share a model call")
}

type digitalHumanDistillChat struct {
	payloads []string
	response string
	err      error
	failAt   int
}

func (m *digitalHumanDistillChat) GetModelID() string   { return "distill-test" }
func (m *digitalHumanDistillChat) GetModelName() string { return "distill-test" }
func (m *digitalHumanDistillChat) ChatStream(context.Context, []chat.Message, *chat.ChatOptions) (<-chan types.StreamResponse, error) {
	return nil, fmt.Errorf("stream is not used for distillation")
}

func (m *digitalHumanDistillChat) Chat(ctx context.Context, messages []chat.Message, _ *chat.ChatOptions) (*types.ChatResponse, error) {
	if _, ok := ctx.Deadline(); !ok {
		return nil, fmt.Errorf("model call must have a timeout")
	}
	m.payloads = append(m.payloads, messages[len(messages)-1].Content)
	if m.failAt > 0 && len(m.payloads) == m.failAt {
		return nil, errors.New("simulated provider interruption")
	}
	return &types.ChatResponse{Content: m.response}, m.err
}

func TestDigitalHumanDistillationDoesNotPublishPartialEvidence(t *testing.T) {
	service, ctx, human, _, _ := setupDigitalHumanWeChat(t)
	scope, err := digitalHumanScope(ctx)
	require.NoError(t, err)
	service.slots = make(chan struct{}, 2)
	persona := map[string]string{}
	for field := range digitalHumanPersonaFields {
		persona[field] = "document evidence"
	}
	encoded, err := json.Marshal(persona)
	require.NoError(t, err)
	model := &digitalHumanDistillChat{response: string(encoded), failAt: 2}
	job := &types.DigitalHumanJob{ID: "partial-job", TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
	require.NoError(t, service.repo.StartJob(ctx, scope, human.ID, human.Revision, job))
	service.slots <- struct{}{}
	service.runDistillation(ctx, scope, job, model, []types.DigitalHumanSource{{Name: "long.txt", Text: strings.Repeat("文", digitalHumanBatchRunes+10)}}, nil, nil, nil, "", human.LatestVersion)
	detail, err := service.repo.Get(ctx, scope, human.ID)
	require.NoError(t, err)
	require.Equal(t, 1, detail.LatestVersion)
	require.Len(t, detail.Versions, 1)
	require.Equal(t, "failed", detail.LastJob.Status)
	require.Equal(t, "clear", detail.Versions[0].Persona["tone"])
}

func TestDigitalHumanIncrementalDistillationRetainsEvidenceAndPreviousVersion(t *testing.T) {
	service, ctx, human, _, _ := setupDigitalHumanWeChat(t)
	scope, err := digitalHumanScope(ctx)
	require.NoError(t, err)
	service.slots = make(chan struct{}, 2)
	previous := []types.DigitalHumanSource{{Name: "old", Text: "already distilled evidence"}}
	added := []types.DigitalHumanSource{{Name: "new", Text: "complete new source, including its ending"}}
	persona := map[string]string{}
	for field := range digitalHumanPersonaFields {
		persona[field] = ""
	}
	persona["narrative_style"] = "new narrative evidence"
	encoded, err := json.Marshal(persona)
	require.NoError(t, err)
	model := &digitalHumanDistillChat{response: string(encoded)}
	job := &types.DigitalHumanJob{ID: "job", TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
	require.NoError(t, service.repo.StartJob(ctx, scope, human.ID, human.Revision, job))
	service.slots <- struct{}{}
	service.runDistillation(ctx, scope, job, model, append(previous, added...), nil, previous, map[string]string{"tone": "prior tone"}, "", human.LatestVersion)
	require.Len(t, model.payloads, 1)
	require.Contains(t, model.payloads[0], added[0].Text)
	require.NotContains(t, model.payloads[0], previous[0].Text)
	detail, err := service.repo.Get(ctx, scope, human.ID)
	require.NoError(t, err)
	require.Equal(t, 2, detail.LatestVersion)
	require.Len(t, detail.Versions, 2)
	version, err := service.repo.Sources(ctx, scope, human.ID, 2)
	require.NoError(t, err)
	require.Equal(t, "prior tone", version.Persona["tone"])
	require.Equal(t, "new narrative evidence", version.Persona["narrative_style"])
	require.Equal(t, append(previous, added...), version.Sources)
	require.Equal(t, "completed", detail.LastJob.Status)
	require.Equal(t, 100, detail.LastJob.Progress)
}

func TestDigitalHumanDistillationResumesOnlyExactFailedBatchSequence(t *testing.T) {
	for _, changedAdvice := range []bool{false, true} {
		t.Run(fmt.Sprintf("changed-advice-%t", changedAdvice), func(t *testing.T) {
			service, ctx, human, _, _ := setupDigitalHumanWeChat(t)
			scope, err := digitalHumanScope(ctx)
			require.NoError(t, err)
			service.slots = make(chan struct{}, 2)
			persona := map[string]string{}
			for field := range digitalHumanPersonaFields {
				persona[field] = "private first-batch finding"
			}
			encoded, err := json.Marshal(persona)
			require.NoError(t, err)
			sources := []types.DigitalHumanSource{{Name: "long.txt", Text: strings.Repeat("文", digitalHumanBatchRunes) + "only unfinished ending"}}
			model := &digitalHumanDistillChat{response: string(encoded), failAt: 2}
			start := func(id string, revision int64) *types.DigitalHumanJob {
				job := &types.DigitalHumanJob{ID: id, TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
				require.NoError(t, service.repo.StartJob(ctx, scope, human.ID, revision, job))
				service.slots <- struct{}{}
				return job
			}
			first := start("first-attempt", human.Revision)
			service.runDistillation(ctx, scope, first, model, sources, nil, nil, nil, "", human.LatestVersion)
			require.Len(t, model.payloads, 2)
			detail, err := service.repo.Get(ctx, scope, human.ID)
			require.NoError(t, err)
			require.Equal(t, 1, detail.LatestVersion)
			require.Equal(t, "failed", detail.LastJob.Status)
			public, err := json.Marshal(detail)
			require.NoError(t, err)
			require.NotContains(t, string(public), "private first-batch finding")
			require.NotContains(t, string(public), "checkpoint")
			// Recreate the service's model/worker state: recovery comes from storage.
			retryModel := &digitalHumanDistillChat{response: string(encoded)}
			advice := ""
			if changedAdvice {
				advice = "use a different narrative style"
			}
			retry := start("retry-attempt", detail.Revision)
			service.runDistillation(ctx, scope, retry, retryModel, sources, nil, nil, nil, advice, human.LatestVersion)
			if changedAdvice {
				require.Len(t, retryModel.payloads, 2, "changed advice must restart all batches")
			} else {
				require.Len(t, retryModel.payloads, 1, "retry must only send the unfinished batch")
				require.NotContains(t, retryModel.payloads[0], strings.Repeat("文", 10))
				require.Contains(t, retryModel.payloads[0], "only unfinished ending")
				require.Contains(t, retryModel.payloads[0], "private first-batch finding")
			}
			detail, err = service.repo.Get(ctx, scope, human.ID)
			require.NoError(t, err)
			require.Equal(t, 2, detail.LatestVersion)
			require.Equal(t, "completed", detail.LastJob.Status)
			version, err := service.repo.Sources(ctx, scope, human.ID, 2)
			require.NoError(t, err)
			require.Equal(t, sources, version.Sources)
		})
	}
}

func TestDigitalHumanCheckpointIdentityIncludesAllGenerationInputs(t *testing.T) {
	scope := types.DigitalHumanScope{TenantID: 7, OwnerID: "owner"}
	sources := []types.DigitalHumanSource{{Name: "one", Text: "first"}, {Name: "two", Text: "second"}}
	persona := map[string]string{"tone": "previous tone"}
	key := digitalHumanDistillFingerprint(scope, "human", "model", "name", 3, persona, sources, "advice")
	keys := []string{
		digitalHumanDistillFingerprint(types.DigitalHumanScope{TenantID: 8, OwnerID: "owner"}, "human", "model", "name", 3, persona, sources, "advice"),
		digitalHumanDistillFingerprint(types.DigitalHumanScope{TenantID: 7, OwnerID: "other"}, "human", "model", "name", 3, persona, sources, "advice"),
		digitalHumanDistillFingerprint(scope, "different-human", "model", "name", 3, persona, sources, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "different-model", "name", 3, persona, sources, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "model", "different-name", 3, persona, sources, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "model", "name", 4, persona, sources, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "model", "name", 3, map[string]string{"tone": "edited"}, sources, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "model", "name", 3, persona, []types.DigitalHumanSource{sources[1], sources[0]}, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "model", "name", 3, persona, []types.DigitalHumanSource{{Name: "one", Text: "edited"}, sources[1]}, "advice"),
		digitalHumanDistillFingerprint(scope, "human", "model", "name", 3, persona, sources, "different advice"),
	}
	for _, different := range keys {
		require.NotEqual(t, key, different)
	}
	require.Equal(t, key, digitalHumanDistillFingerprint(scope, "human", "model", "name", 3, persona, sources, "advice"))
}

func TestDigitalHumanPersonaRejectsIncompleteOrEmptyModelOutput(t *testing.T) {
	persona := map[string]string{}
	for field := range digitalHumanPersonaFields {
		persona[field] = ""
	}
	encoded, err := json.Marshal(persona)
	require.NoError(t, err)
	_, err = parseDigitalHumanPersona(string(encoded))
	require.Error(t, err)
	persona["tone"] = "Evidence from the source"
	encoded, err = json.Marshal(persona)
	require.NoError(t, err)
	parsed, err := parseDigitalHumanPersona("```json\n" + string(encoded) + "\n```")
	require.NoError(t, err)
	require.Equal(t, persona, parsed)
	delete(persona, "identity")
	encoded, err = json.Marshal(persona)
	require.NoError(t, err)
	_, err = parseDigitalHumanPersona(string(encoded))
	require.Error(t, err)
	_, err = parseDigitalHumanPersona(`{"tone": ["invalid"]}`)
	require.Error(t, err)
}