xinmi/internal/im/agent_lifecycle.go

package im

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/Tencent/WeKnora/internal/types"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type verifiedAgentExecutionKey struct{}
type verifiedAgentExecution struct {
	TenantID                 uint64
	AgentID, OwnerID, ChannelID string
	PrincipalID              string
}

// IsVerifiedAgentExecution accepts only a proof minted inside the IM runtime
// after reading an enabled channel and its exact Agent from durable storage.
// A synthetic user, IM principal or request parameter alone never grants it.
func IsVerifiedAgentExecution(ctx context.Context, agent *types.CustomAgent) bool {
	proof, ok := ctx.Value(verifiedAgentExecutionKey{}).(verifiedAgentExecution)
	if !ok || agent == nil || proof.ChannelID == "" || proof.AgentID != agent.ID || proof.TenantID != agent.TenantID || proof.OwnerID != agent.CreatedBy {
		return false
	}
	tenant, hasTenant := types.TenantIDFromContext(ctx)
	user, hasUser := types.UserIDFromContext(ctx)
	principal, hasPrincipal := types.PrincipalFromContext(ctx)
	return hasTenant && hasUser && tenant == proof.TenantID && user == fmt.Sprintf("system-%d", tenant) &&
		hasPrincipal && principal.Type == types.PrincipalIMUser && principal.ID == proof.PrincipalID
}

func (s *Service) resolveConfiguredChannelExecution(ctx context.Context, channel *IMChannel) (context.Context, *types.CustomAgent, error) {
	agent, err := s.resolveConfiguredChannelAgent(ctx, channel)
	if err != nil {
		return ctx, nil, err
	}
	principal, hasPrincipal := types.PrincipalFromContext(ctx)
	tenant, hasTenant := types.TenantIDFromContext(ctx)
	user, hasUser := types.UserIDFromContext(ctx)
	if !hasPrincipal || principal.Type != types.PrincipalIMUser || !hasTenant || tenant != channel.TenantID || !hasUser || user != fmt.Sprintf("system-%d", tenant) {
		return ctx, nil, fmt.Errorf("verified IM identity is required for channel execution")
	}
	proof := verifiedAgentExecution{TenantID: channel.TenantID, AgentID: agent.ID, OwnerID: agent.CreatedBy, ChannelID: channel.ID, PrincipalID: principal.ID}
	return context.WithValue(ctx, verifiedAgentExecutionKey{}, proof), agent, nil
}

// InvalidateDeletedChannels is called only after durable deletion commits.
func (s *Service) InvalidateDeletedChannels(ids []string) {
	for _, id := range ids {
		s.StopChannel(id)
		s.publishChannelConfigChange(id)
	}
}

func (s *Service) cancelChannelRequests(id string) {
	s.inflight.Range(func(key, value any) bool {
		userKey, ok := key.(string)
		if ok && strings.HasPrefix(userKey, id+":") {
			if entry, ok := value.(*inflightEntry); ok && entry.cancel != nil {
				entry.cancel()
			}
		}
		return true
	})
}

// Serialize authoring with agent deletion so a concurrent channel create or
// reassignment cannot leave a new channel pointing at a deleted custom agent.
func lockChannelAgent(tx *gorm.DB, channel *IMChannel) error {
	if types.IsBuiltinAgentID(channel.AgentID) {
		return nil // Built-in agents cannot be deleted.
	}
	var agent types.CustomAgent
	return tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id").
		Where("id = ? AND tenant_id = ?", channel.AgentID, channel.TenantID).First(&agent).Error
}

// resolveConfiguredChannelAgent is a trusted runtime path, not a user Agent
// lookup. A persisted, enabled channel is the explicit authorization to execute
// exactly its configured Agent, including an owner's private digital human.
func (s *Service) resolveConfiguredChannelAgent(ctx context.Context, channel *IMChannel) (*types.CustomAgent, error) {
	if channel == nil || channel.ID == "" || channel.AgentID == "" {
		return nil, fmt.Errorf("channel has no configured agent")
	}
	var fresh IMChannel
	if err := s.db.WithContext(ctx).Where("id = ? AND tenant_id = ? AND agent_id = ? AND enabled = ?", channel.ID, channel.TenantID, channel.AgentID, true).First(&fresh).Error; err != nil {
		return nil, err
	}
	var agent types.CustomAgent
	err := s.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", fresh.AgentID, fresh.TenantID).First(&agent).Error
	if errors.Is(err, gorm.ErrRecordNotFound) && types.IsBuiltinAgentID(fresh.AgentID) {
		return s.agentService.GetAgentByID(ctx, fresh.AgentID)
	}
	if err != nil {
		return nil, err
	}
	agent.EnsureDefaults()
	return &agent, nil
}