From 1c80f3e52f8acda073a8e5ca2d7f8ad858ad8a42 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Thu, 4 Sep 2025 06:36:25 +0800
Subject: [PATCH] feat(agent): add agent iteration and count functions;
 refactor Forward method to return *http.Response

---
 agent/pkg/agent/agent_pool.go    | 10 ++++++++++
 agent/pkg/agent/http_requests.go |  8 +++-----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/agent/pkg/agent/agent_pool.go b/agent/pkg/agent/agent_pool.go
index 49c3e6cd..a22b8d31 100644
--- a/agent/pkg/agent/agent_pool.go
+++ b/agent/pkg/agent/agent_pool.go
@@ -1,6 +1,8 @@
 package agent
 
 import (
+	"iter"
+
 	"github.com/yusing/go-proxy/internal/common"
 	"github.com/yusing/go-proxy/internal/utils/functional"
 )
@@ -51,6 +53,14 @@ func ListAgents() []*AgentConfig {
 	return agents
 }
 
+func IterAgents() iter.Seq2[string, *AgentConfig] {
+	return agentPool.Range
+}
+
+func NumAgents() int {
+	return agentPool.Size()
+}
+
 func getAgentByAddr(addr string) (agent *AgentConfig, ok bool) {
 	agent, ok = agentPool.Load(addr)
 	return
diff --git a/agent/pkg/agent/http_requests.go b/agent/pkg/agent/http_requests.go
index 6dc655d9..3bb2afca 100644
--- a/agent/pkg/agent/http_requests.go
+++ b/agent/pkg/agent/http_requests.go
@@ -16,7 +16,7 @@ func (cfg *AgentConfig) Do(ctx context.Context, method, endpoint string, body io
 	return cfg.httpClient.Do(req)
 }
 
-func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) ([]byte, int, error) {
+func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) (*http.Response, error) {
 	req = req.WithContext(req.Context())
 	req.URL.Host = AgentHost
 	req.URL.Scheme = "https"
@@ -24,11 +24,9 @@ func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) ([]byte, int
 	req.RequestURI = ""
 	resp, err := cfg.httpClient.Do(req)
 	if err != nil {
-		return nil, 0, err
+		return nil, err
 	}
-	defer resp.Body.Close()
-	data, _ := io.ReadAll(resp.Body)
-	return data, resp.StatusCode, nil
+	return resp, nil
 }
 
 func (cfg *AgentConfig) Fetch(ctx context.Context, endpoint string) ([]byte, int, error) {