dify-plugin-daemon/internal/core/session_manager/session.go

package session_manager

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation"
	"github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types"
	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
	"github.com/langgenius/dify-plugin-daemon/pkg/utils/cache"
	"github.com/langgenius/dify-plugin-daemon/pkg/utils/log"
	"github.com/langgenius/dify-plugin-daemon/pkg/utils/parser"
)

var (
	sessions    map[string]*Session = map[string]*Session{}
	sessionLock sync.RWMutex
)

// session need to implement the backwards_invocation.BackwardsInvocationWriter interface
type Session struct {
	ID                  string                                          `json:"id"`
	TraceContext        log.TraceContext                                `json:"trace_context"`
	IdentityContext     log.Identity                                    `json:"identity_context"`
	runtime             plugin_entities.PluginRuntimeSessionIOInterface `json:"-"`
	backwardsInvocation dify_invocation.BackwardsInvocation             `json:"-"`
	requestContext      context.Context                                 `json:"-"`

	TenantID               string                                 `json:"tenant_id"`
	UserID                 string                                 `json:"user_id"`
	PluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier `json:"plugin_unique_identifier"`
	ClusterID              string                                 `json:"cluster_id"`
	InvokeFrom             access_types.PluginAccessType          `json:"invoke_from"`
	Action                 access_types.PluginAccessAction        `json:"action"`
	Declaration            *plugin_entities.PluginDeclaration     `json:"declaration"`

	// information about incoming request
	ConversationID *string        `json:"conversation_id"`
	MessageID      *string        `json:"message_id"`
	AppID          *string        `json:"app_id"`
	EndpointID     *string        `json:"endpoint_id"`
	Context        map[string]any `json:"context"`
}

func sessionKey(id string) string {
	return fmt.Sprintf("session_info:%s", id)
}

type NewSessionPayload struct {
	TenantID               string                                 `json:"tenant_id"`
	UserID                 string                                 `json:"user_id"`
	PluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier `json:"plugin_unique_identifier"`
	ClusterID              string                                 `json:"cluster_id"`
	InvokeFrom             access_types.PluginAccessType          `json:"invoke_from"`
	Action                 access_types.PluginAccessAction        `json:"action"`
	Declaration            *plugin_entities.PluginDeclaration     `json:"declaration"`
	BackwardsInvocation    dify_invocation.BackwardsInvocation    `json:"backwards_invocation"`
	IgnoreCache            bool                                   `json:"ignore_cache"`
	ConversationID         *string                                `json:"conversation_id"`
	MessageID              *string                                `json:"message_id"`
	AppID                  *string                                `json:"app_id"`
	EndpointID             *string                                `json:"endpoint_id"`
	Context                map[string]any                         `json:"context"`
	RequestContext         context.Context                        `json:"-"`
}

func NewSession(payload NewSessionPayload) *Session {
	s := &Session{
		ID:                     uuid.New().String(),
		TenantID:               payload.TenantID,
		UserID:                 payload.UserID,
		PluginUniqueIdentifier: payload.PluginUniqueIdentifier,
		ClusterID:              payload.ClusterID,
		InvokeFrom:             payload.InvokeFrom,
		Action:                 payload.Action,
		Declaration:            payload.Declaration,
		backwardsInvocation:    payload.BackwardsInvocation,
		requestContext:         payload.RequestContext,
		ConversationID:         payload.ConversationID,
		MessageID:              payload.MessageID,
		AppID:                  payload.AppID,
		EndpointID:             payload.EndpointID,
		Context:                payload.Context,
	}

	s.propagateTraceContext()

	sessionLock.Lock()
	sessions[s.ID] = s
	sessionLock.Unlock()

	if !payload.IgnoreCache {
		if err := cache.Store(sessionKey(s.ID), s, time.Minute*30); err != nil {
			log.ErrorContext(
				s.RequestContext(),
				"set session info to cache failed",
				"session_id", s.ID,
				"cache_key", sessionKey(s.ID),
				"cluster_id", s.ClusterID,
				"plugin_unique_identifier", s.PluginUniqueIdentifier.String(),
				"error", err,
			)
		}
	}

	return s
}

// GetSession tries to get the session from local memory first, then from cache
//
// Usally, `GetSession` is only used in serverless mode, at this scenario, GetSession
// try to get trace context from cache and reload it to session
func GetSession(id string) (*Session, error) {
	sessionLock.RLock()
	session := sessions[id]
	sessionLock.RUnlock()

	if session == nil {
		// if session not found, it may be generated by another node, try to get it from cache
		session, err := cache.Get[Session](sessionKey(id))
		if err != nil {
			return nil, errors.Join(err, errors.New("failed to get session info from cache"))
		}

		// get a background context
		ctx := context.Background()
		// initalize trace context
		ctx = log.WithTrace(ctx, session.TraceContext)
		ctx = log.WithIdentity(ctx, session.IdentityContext)

		session.requestContext = ctx
		session.propagateTraceContext()
		return session, nil
	}

	return session, nil
}

type DeleteSessionPayload struct {
	ID          string `json:"id"`
	IgnoreCache bool   `json:"ignore_cache"`
}

func DeleteSession(payload DeleteSessionPayload) {
	sessionLock.Lock()
	session := sessions[payload.ID]
	delete(sessions, payload.ID)
	sessionLock.Unlock()

	if !payload.IgnoreCache {
		cacheKey := sessionKey(payload.ID)
		if _, err := cache.Del(cacheKey); err != nil {
			if session != nil {
				log.ErrorContext(
					session.RequestContext(),
					"delete session info from cache failed",
					"session_id", payload.ID,
					"cache_key", cacheKey,
					"cluster_id", session.ClusterID,
					"plugin_unique_identifier", session.PluginUniqueIdentifier.String(),
					"error", err,
				)
			} else {
				log.ErrorContext(
					context.Background(),
					"delete session info from cache failed",
					"session_id", payload.ID,
					"cache_key", cacheKey,
					"error", err,
				)
			}
		}
	}
}

type CloseSessionPayload struct {
	IgnoreCache bool `json:"ignore_cache"`
}

func (s *Session) Close(payload CloseSessionPayload) {
	DeleteSession(DeleteSessionPayload{
		ID:          s.ID,
		IgnoreCache: payload.IgnoreCache,
	})
}

func (s *Session) BindRuntime(runtime plugin_entities.PluginRuntimeSessionIOInterface) {
	s.runtime = runtime
}

func (s *Session) Runtime() plugin_entities.PluginRuntimeSessionIOInterface {
	return s.runtime
}

func (s *Session) BindBackwardsInvocation(backwardsInvocation dify_invocation.BackwardsInvocation) {
	s.backwardsInvocation = backwardsInvocation
	s.propagateTraceContext()
}

func (s *Session) BackwardsInvocation() dify_invocation.BackwardsInvocation {
	return s.backwardsInvocation
}

// propagateTraceContext propagates the trace context from request context
// to session, so that some of the downstream calls may marshall the trace context
func (s *Session) propagateTraceContext() {
	if s.backwardsInvocation != nil && s.requestContext != nil {
		s.backwardsInvocation.SetContext(s.requestContext)
	}

	s.TraceContext, _ = log.TraceFromContext(s.requestContext)
	s.IdentityContext, _ = log.IdentityFromContext(s.requestContext)
}

func (s *Session) RequestContext() context.Context {
	if s.requestContext == nil {
		return context.Background()
	}
	return s.requestContext
}

type PLUGIN_IN_STREAM_EVENT string

const (
	PLUGIN_IN_STREAM_EVENT_REQUEST  PLUGIN_IN_STREAM_EVENT = "request"
	PLUGIN_IN_STREAM_EVENT_RESPONSE PLUGIN_IN_STREAM_EVENT = "backwards_response"
)

func (s *Session) Message(event PLUGIN_IN_STREAM_EVENT, data any) []byte {
	return parser.MarshalJsonBytes(map[string]any{
		"session_id":      s.ID,
		"conversation_id": s.ConversationID,
		"message_id":      s.MessageID,
		"app_id":          s.AppID,
		"endpoint_id":     s.EndpointID,
		"context":         s.Context,
		"event":           event,
		"data":            data,
	})
}

func (s *Session) Write(event PLUGIN_IN_STREAM_EVENT, action access_types.PluginAccessAction, data any) error {
	if s.runtime == nil {
		return errors.New("runtime not bound")
	}
	return s.runtime.Write(s.ID, action, s.Message(event, data))
}