io.go 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. package serverless_runtime
  2. import (
  3. "bufio"
  4. "bytes"
  5. "context"
  6. "fmt"
  7. "net/http"
  8. "net/url"
  9. "time"
  10. "github.com/langgenius/dify-plugin-daemon/internal/core/plugin_daemon/access_types"
  11. "github.com/langgenius/dify-plugin-daemon/internal/utils/log"
  12. "github.com/langgenius/dify-plugin-daemon/internal/utils/parser"
  13. "github.com/langgenius/dify-plugin-daemon/internal/utils/routine"
  14. "github.com/langgenius/dify-plugin-daemon/pkg/entities"
  15. "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
  16. )
  17. func (r *AWSPluginRuntime) Listen(sessionId string) *entities.Broadcast[plugin_entities.SessionMessage] {
  18. l := entities.NewBroadcast[plugin_entities.SessionMessage]()
  19. // store the listener
  20. r.listeners.Store(sessionId, l)
  21. return l
  22. }
  23. // For AWS Lambda, write is equivalent to http request, it's not a normal stream like stdio and tcp
  24. func (r *AWSPluginRuntime) Write(sessionId string, action access_types.PluginAccessAction, data []byte) {
  25. l, ok := r.listeners.Load(sessionId)
  26. if !ok {
  27. log.Error("session %s not found", sessionId)
  28. return
  29. }
  30. url, err := url.JoinPath(r.LambdaURL, "invoke")
  31. if err != nil {
  32. l.Send(plugin_entities.SessionMessage{
  33. Type: plugin_entities.SESSION_MESSAGE_TYPE_ERROR,
  34. Data: parser.MarshalJsonBytes(plugin_entities.ErrorResponse{
  35. ErrorType: "PluginDaemonInnerError",
  36. Message: fmt.Sprintf("Error creating request: %v", err),
  37. }),
  38. })
  39. l.Close()
  40. r.Error(fmt.Sprintf("Error creating request: %v", err))
  41. return
  42. }
  43. url += "?action=" + string(action)
  44. connectTime := 240 * time.Second
  45. // create a new http request
  46. ctx, cancel := context.WithTimeout(context.Background(), connectTime)
  47. time.AfterFunc(connectTime, cancel)
  48. req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(data))
  49. if err != nil {
  50. r.Error(fmt.Sprintf("Error creating request: %v", err))
  51. return
  52. }
  53. req.Header.Set("Content-Type", "application/json")
  54. req.Header.Set("Accept", "text/event-stream")
  55. req.Header.Set("Dify-Plugin-Session-ID", sessionId)
  56. routine.Submit(map[string]string{
  57. "module": "serverless_runtime",
  58. "function": "Write",
  59. "session_id": sessionId,
  60. "lambda_url": r.LambdaURL,
  61. }, func() {
  62. // remove the session from listeners
  63. defer r.listeners.Delete(sessionId)
  64. defer l.Close()
  65. defer l.Send(plugin_entities.SessionMessage{
  66. Type: plugin_entities.SESSION_MESSAGE_TYPE_END,
  67. Data: []byte(""),
  68. })
  69. response, err := r.client.Do(req)
  70. if err != nil {
  71. l.Send(plugin_entities.SessionMessage{
  72. Type: plugin_entities.SESSION_MESSAGE_TYPE_ERROR,
  73. Data: parser.MarshalJsonBytes(plugin_entities.ErrorResponse{
  74. ErrorType: "PluginDaemonInnerError",
  75. Message: fmt.Sprintf("Error sending request to aws lambda: %v", err),
  76. }),
  77. })
  78. r.Error(fmt.Sprintf("Error sending request to aws lambda: %v", err))
  79. return
  80. }
  81. // write to data stream
  82. scanner := bufio.NewScanner(response.Body)
  83. defer response.Body.Close()
  84. // TODO: set a reasonable buffer size or use a reader, this is a temporary solution
  85. scanner.Buffer(make([]byte, 1024), 5*1024*1024)
  86. sessionAlive := true
  87. for scanner.Scan() && sessionAlive {
  88. bytes := scanner.Bytes()
  89. if len(bytes) == 0 {
  90. continue
  91. }
  92. plugin_entities.ParsePluginUniversalEvent(
  93. bytes,
  94. response.Status,
  95. func(session_id string, data []byte) {
  96. sessionMessage, err := parser.UnmarshalJsonBytes[plugin_entities.SessionMessage](data)
  97. if err != nil {
  98. l.Send(plugin_entities.SessionMessage{
  99. Type: plugin_entities.SESSION_MESSAGE_TYPE_ERROR,
  100. Data: parser.MarshalJsonBytes(plugin_entities.ErrorResponse{
  101. ErrorType: "PluginDaemonInnerError",
  102. Message: fmt.Sprintf("failed to parse session message %s, err: %v", bytes, err),
  103. }),
  104. })
  105. sessionAlive = false
  106. }
  107. l.Send(sessionMessage)
  108. },
  109. func() {},
  110. func(err string) {
  111. l.Send(plugin_entities.SessionMessage{
  112. Type: plugin_entities.SESSION_MESSAGE_TYPE_ERROR,
  113. Data: parser.MarshalJsonBytes(plugin_entities.ErrorResponse{
  114. ErrorType: "PluginDaemonInnerError",
  115. Message: fmt.Sprintf("encountered an error: %v", err),
  116. }),
  117. })
  118. },
  119. func(message string) {},
  120. )
  121. }
  122. if err := scanner.Err(); err != nil {
  123. l.Send(plugin_entities.SessionMessage{
  124. Type: plugin_entities.SESSION_MESSAGE_TYPE_ERROR,
  125. Data: parser.MarshalJsonBytes(plugin_entities.ErrorResponse{
  126. ErrorType: "PluginDaemonInnerError",
  127. Message: fmt.Sprintf("failed to read response body: %v", err),
  128. }),
  129. })
  130. }
  131. })
  132. }