feat: add retry mechanism for serverless invocation on 502 errors (#569)

* feat: add retry mechanism for serverless invocation on 502 errors

- Add MAX_SERVERLESS_RETRY_TIMES config (default: 3) for configurable retry attempts
- Implement exponential backoff retry logic (500ms, 1s, 2s, ...) for 502 Bad Gateway errors
- Only retry on 502 status code as it indicates transient AWS Lambda gateway issues
- Add comprehensive test suite with 11 test cases covering all retry scenarios
- Ensure proper HTTP response body cleanup on retries to prevent resource leaks

* refactor: simplify retry logic and error handling in serverless invocation

* fix: max time

* fix: test

Author: Yeuoly

internal/core/serverless_runtime/io.go

@@ -6,7 +6,9 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"net/url"
+	"time"
 
 	"github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types"
 	"github.com/langgenius/dify-plugin-daemon/pkg/entities"
@@ -27,6 +29,86 @@ func (r *ServerlessPluginRuntime) Listen(sessionId string) (
 	return l, nil
 }
 
+// shouldRetryStatusCode checks if the HTTP status code warrants a retry
+// Only 502 (Bad Gateway) errors are retried as they typically indicate temporary gateway issues
+//
+// To some AWS Lambda gateway errors, 502 randomly happens, and it's usually transient.
+// Thus we implement a retry mechanism for 502 errors.
+func shouldRetryStatusCode(statusCode int) bool {
+	return statusCode == 502
+}
+
+// invokeServerlessWithRetry invokes the serverless endpoint with retry logic
+// It will retry up to MaxRetryTimes attempts on 502 errors with exponential backoff
+// Backoff duration is capped at 30 seconds to prevent unreasonable wait times
+func (r *ServerlessPluginRuntime) invokeServerlessWithRetry(
+	url string,
+	sessionId string,
+	data []byte,
+) (*http.Response, error) {
+	const maxBackoffDuration = 30 * time.Second
+
+	var lastErr error
+
+	maxRetries := r.MaxRetryTimes
+	if maxRetries <= 0 {
+		maxRetries = 1
+	}
+
+	for attempt := 0; attempt < maxRetries; attempt++ {
+		// Apply exponential backoff for retry attempts (500ms, 1000ms, 2000ms, ...)
+		// Capped at 30 seconds to prevent unreasonable wait times
+		if attempt > 0 {
+			backoffDuration := time.Duration(500*(1<<uint(attempt-1))) * time.Millisecond
+			if backoffDuration > maxBackoffDuration {
+				backoffDuration = maxBackoffDuration
+			}
+			time.Sleep(backoffDuration)
+		}
+
+		// Make HTTP request to serverless endpoint
+		response, err := http_requests.Request(
+			r.Client, url, "POST",
+			http_requests.HttpHeader(map[string]string{
+				"Content-Type":           "application/json",
+				"Accept":                 "text/event-stream",
+				"Dify-Plugin-Session-ID": sessionId,
+			}),
+			http_requests.HttpPayloadReader(io.NopCloser(bytes.NewReader(data))),
+			http_requests.HttpReadTimeout(int64(r.PluginMaxExecutionTimeout*1000)),
+		)
+
+		if err != nil {
+			lastErr = fmt.Errorf("attempt %d/%d failed: %w", attempt+1, maxRetries, err)
+			continue
+		}
+
+		statusCode := response.StatusCode
+		// Success - return immediately
+		if statusCode >= 200 && statusCode < 300 {
+			return response, nil
+		}
+
+		// Check if status code should trigger a retry (502 Bad Gateway only)
+		if shouldRetryStatusCode(statusCode) {
+			if response.Body != nil {
+				response.Body.Close()
+			}
+			lastErr = fmt.Errorf("attempt %d/%d failed with status code: %d", attempt+1, maxRetries, statusCode)
+			continue
+		}
+
+		// Non-retryable error - return immediately
+		return response, nil
+	}
+
+	if lastErr != nil {
+		return nil, fmt.Errorf("all %d attempts failed, last error: %w", maxRetries, lastErr)
+	}
+
+	return nil, fmt.Errorf("all %d attempts failed with unknown error", maxRetries)
+}
+
 // For Serverless, write is equivalent to http request, it's not a normal stream like stdio and tcp
 func (r *ServerlessPluginRuntime) Write(
 	sessionId string,
@@ -49,26 +131,15 @@ func (r *ServerlessPluginRuntime) Write(
 		routinepkg.RoutineLabelKeySessionID: sessionId,
 		routinepkg.RoutineLabelKeyLambdaURL: r.LambdaURL,
 	}, func() {
-		// remove the session from listeners
 		defer r.listeners.Delete(sessionId)
 		defer l.Close()
 		defer l.Send(plugin_entities.SessionMessage{
 			Type: plugin_entities.SESSION_MESSAGE_TYPE_END,
 			Data: []byte(""),
 		})
 
-		// create a new http request to serverless runtimes
 		url += "?action=" + string(action)
-		response, err := http_requests.Request(
-			r.Client, url, "POST",
-			http_requests.HttpHeader(map[string]string{
-				"Content-Type":           "application/json",
-				"Accept":                 "text/event-stream",
-				"Dify-Plugin-Session-ID": sessionId,
-			}),
-			http_requests.HttpPayloadReader(io.NopCloser(bytes.NewReader(data))),
-			http_requests.HttpReadTimeout(int64(r.PluginMaxExecutionTimeout*1000)),
-		)
+		response, err := r.invokeServerlessWithRetry(url, sessionId, data)
 		if err != nil {
 			l.Send(plugin_entities.SessionMessage{
 				Type: plugin_entities.SESSION_MESSAGE_TYPE_ERROR,
@@ -80,7 +151,6 @@ func (r *ServerlessPluginRuntime) Write(
 			return
 		}
 
-		// write to data stream
 		scanner := bufio.NewScanner(response.Body)
 		defer response.Body.Close()