xinmi/internal/im/agent_lifecycle_test.go

package im

import (
	"context"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func TestDigitalHumanIMConfiguredAgentIsScopedAndNeverFallsBack(t *testing.T) {
	db := newLifecycleTestDB(t)
	channel := createLifecycleChannel(t, db, "private-channel", "private-agent")
	var agent types.CustomAgent
	require.NoError(t, db.First(&agent, "id = ?", channel.AgentID).Error)
	agent.CreatedBy = "owner"
	agent.Config.SystemPrompt = "<!-- weknora:digital-human:private:start -->Private persona"
	require.NoError(t, db.Save(&agent).Error)
	svc := newLifecycleTestService(db, nil, "one")
	t.Cleanup(svc.Stop)
	ctx := withIMIdentity(context.Background(), 1, channel.ID, &IncomingMessage{UserID: "external-user"})
	resolved, err := svc.resolveConfiguredChannelAgent(ctx, channel)
	require.NoError(t, err)
	require.Contains(t, resolved.Config.SystemPrompt, "Private persona")
	foreign := *channel
	foreign.TenantID = 2
	_, err = svc.resolveConfiguredChannelAgent(ctx, &foreign)
	require.ErrorIs(t, err, gorm.ErrRecordNotFound)
	foreign = *channel
	foreign.AgentID = "unconfigured-agent"
	_, err = svc.resolveConfiguredChannelAgent(ctx, &foreign)
	require.ErrorIs(t, err, gorm.ErrRecordNotFound)
	require.NoError(t, db.Delete(&agent).Error)
	_, err = svc.resolveConfiguredChannelAgent(ctx, channel)
	require.ErrorIs(t, err, gorm.ErrRecordNotFound)
}

func TestDigitalHumanIMLateStartupCannotOutliveDeletion(t *testing.T) {
	db := newLifecycleTestDB(t)
	channel := createLifecycleChannel(t, db, "late-channel", "agent")
	svc := newLifecycleTestService(db, nil, "one")
	t.Cleanup(svc.Stop)
	entered, proceed, stopped := make(chan struct{}), make(chan struct{}), make(chan struct{})
	svc.RegisterAdapterFactory("test", func(context.Context, *IMChannel, func(context.Context, *IncomingMessage) error) (Adapter, context.CancelFunc, error) {
		close(entered)
		<-proceed
		return &lifecycleTestAdapter{}, func() { close(stopped) }, nil
	})
	started := make(chan error, 1)
	go func() { started <- svc.StartChannel(channel) }()
	select {
	case <-entered:
	case <-time.After(time.Second):
		t.Fatal("factory did not start")
	}
	require.NoError(t, db.Delete(channel).Error)
	svc.InvalidateDeletedChannels([]string{channel.ID})
	close(proceed)
	require.ErrorIs(t, <-started, gorm.ErrRecordNotFound)
	select {
	case <-stopped:
	default:
		t.Fatal("late adapter was not stopped")
	}
	_, _, active := svc.GetChannelAdapter(channel.ID)
	require.False(t, active)
}

func TestDigitalHumanIMDeletedChannelStopsRuntimeAndRejectsStaleMessages(t *testing.T) {
	db := newLifecycleTestDB(t)
	channel := createLifecycleChannel(t, db, "deleted-channel", "agent")
	svc := newLifecycleTestService(db, nil, "one")
	counters := &lifecycleFactoryCounters{}
	svc.RegisterAdapterFactory("test", counters.factory())
	t.Cleanup(svc.Stop)
	require.NoError(t, svc.StartChannel(channel))
	requestCtx, cancel := context.WithCancel(context.Background())
	defer cancel()
	svc.inflight.Store(channel.ID+":user:chat", &inflightEntry{cancel: cancel})
	otherCtx, otherCancel := context.WithCancel(context.Background())
	defer otherCancel()
	svc.inflight.Store("other:user:chat", &inflightEntry{cancel: otherCancel})
	require.NoError(t, db.Delete(channel).Error)
	// A missed invalidation event must still reject the next long-poll message.
	require.Error(t, svc.HandleMessage(context.Background(), &IncomingMessage{Content: "hello"}, channel.ID))
	require.Equal(t, int32(1), counters.stops.Load())
	require.ErrorIs(t, requestCtx.Err(), context.Canceled)
	require.NoError(t, otherCtx.Err())
	_, _, active := svc.GetChannelAdapter(channel.ID)
	require.False(t, active)
	svc.InvalidateDeletedChannels([]string{channel.ID})
	require.Equal(t, int32(1), counters.stops.Load())
}

func TestDigitalHumanIMChannelWritesCannotRestoreDeletedBindings(t *testing.T) {
	db := newLifecycleTestDB(t)
	channel := createLifecycleChannel(t, db, "channel", "agent")
	svc := newLifecycleTestService(db, nil, "one")
	t.Cleanup(svc.Stop)
	stale := *channel
	require.NoError(t, db.Delete(channel).Error)
	require.ErrorIs(t, svc.UpdateChannel(&stale), gorm.ErrRecordNotFound)
	var active int64
	require.NoError(t, db.Model(&IMChannel{}).Where("id = ?", channel.ID).Count(&active).Error)
	require.Zero(t, active)
	require.NoError(t, db.Where("id = ?", "agent").Delete(&types.CustomAgent{}).Error)
	stale.ID = "new-channel"
	require.ErrorIs(t, svc.CreateChannel(&stale), gorm.ErrRecordNotFound)
	require.NoError(t, db.Model(&IMChannel{}).Count(&active).Error)
	require.Zero(t, active)
}

func TestDigitalHumanIMDisabledChannelRemainsDisabledAtCreation(t *testing.T) {
	db := newLifecycleTestDB(t)
	existing := createLifecycleChannel(t, db, "existing-channel", "agent")
	channel := &IMChannel{ID: "disabled-channel", TenantID: existing.TenantID, AgentID: existing.AgentID,
		Platform: "test", Name: "Disabled connection", Enabled: false, Credentials: types.JSON("{}")}
	svc := newLifecycleTestService(db, nil, "one")
	t.Cleanup(svc.Stop)
	counters := &lifecycleFactoryCounters{}
	svc.RegisterAdapterFactory("test", counters.factory())
	require.NoError(t, svc.CreateChannel(channel))
	require.False(t, channel.Enabled)
	var saved IMChannel
	require.NoError(t, db.First(&saved, "id = ?", channel.ID).Error)
	require.False(t, saved.Enabled)
	_, _, active := svc.GetChannelAdapter(channel.ID)
	require.False(t, active)
}