perf: 添加任务中断

jumpserver · Dec 12, 2023 · 010e583 · 010e583
1 parent f36cafb
commit 010e583
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/pkg/httpd/router/handler.go b/pkg/httpd/router/handler.go
@@ -20,6 +20,15 @@ func GetSession(sessionID string) (*jms.JMSSession, error) {
 	return nil, errors.New("not found conversation")
 }
 
+
+func GetSystemSession(sessionID string) (*jms.JMSSystemSession, error) {
+	systemJmsSession, _ := jms.GlobalSystemSessionManager.GetSystemSession(sessionID)
+	if systemJmsSession != nil{
+		return systemJmsSession, nil
+	}
+	return nil, errors.New("not found conversation")
+}
+
 func (s *_HandlerApi) InterruptCurrentAskHandler(ctx *gin.Context) {
 	var conversation schemas.Conversation
 	if err := ctx.ShouldBindJSON(&conversation); err != nil {
@@ -28,11 +37,17 @@ func (s *_HandlerApi) InterruptCurrentAskHandler(ctx *gin.Context) {
 	}
 
 	jmsSession, err := GetSession(conversation.ID)
-	if err != nil {
+	jmsSystemSession, systemErr := GetSystemSession(conversation.ID)
+	if err != nil && systemErr != nil{
 		ctx.JSON(http.StatusNotFound, gin.H{"message": err.Error()})
 		return
 	}
-	jmsSession.CurrentAskInterrupt = true
+
+	if err == nil {
+		jmsSession.CurrentAskInterrupt = true
+	} else if systemErr == nil {
+		jmsSystemSession.CurrentAskInterrupt = true
+	}
 	ctx.JSON(http.StatusOK, gin.H{"message": "Success"})
 }
 

diff --git a/pkg/jms/session_manager.go b/pkg/jms/session_manager.go
@@ -8,7 +8,6 @@ import (
 type Session interface {
 	GetID() string
 }
-
 type SessionManager struct {
 	store sync.Map
 }