socks5 support

main.go

@@ -15,6 +15,7 @@ import (
   "io"
   "net"
   "net/http"
+  "net/url"
   "os"
   "os/exec"
   "os/signal"
@@ -61,9 +62,24 @@ func internalKey(k string) bool {
   return false
 }
 
+func proxyFromEnv(req *http.Request) (*url.URL, error) {
+  socks := strings.TrimSpace(os.Getenv("SOCKS_PROXY"))
+  if socks == "" {
+    socks = strings.TrimSpace(os.Getenv("socks_proxy"))
+  }
+  if socks != "" {
+    if !strings.Contains(socks, "://") {
+      socks = "socks5://" + socks
+    }
+    return url.Parse(socks)
+  }
+  return http.ProxyFromEnvironment(req)
+}
+
 // llmTransport is a shared HTTP transport reused across all LLM calls so that
 // connections are pooled instead of recreated per request.
 var llmTransport = &http.Transport{
+  Proxy:       proxyFromEnv,
   DialContext: (&net.Dialer{Timeout: 300 * time.Second}).DialContext,
 }
 
@@ -77,7 +93,13 @@ func atoiD(s string, d int) int {
 }
 
 func queryModelsContextWindow(cfg *Cfg) int {
-  client := &http.Client{Timeout: 3 * time.Second}
+  client := &http.Client{
+    Transport: &http.Transport{
+      Proxy:       proxyFromEnv,
+      DialContext: (&net.Dialer{Timeout: 3 * time.Second}).DialContext,
+    },
+    Timeout: 3 * time.Second,
+  }
   req, err := http.NewRequest("GET", strings.TrimRight(cfg.Endpoint, "/")+"/models", nil)
   if err != nil { return 0 }
   req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)")
@@ -143,7 +165,13 @@ func fetchContextWindow(cfg *Cfg) int {
 func listModels(cfg *Cfg) (string, error) {
   t := cfg.Timeout
   if t < 10 { t = 10 }
-  client := &http.Client{Timeout: time.Duration(t) * time.Second}
+  client := &http.Client{
+    Transport: &http.Transport{
+      Proxy:       proxyFromEnv,
+      DialContext: (&net.Dialer{Timeout: time.Duration(t) * time.Second}).DialContext,
+    },
+    Timeout: time.Duration(t) * time.Second,
+  }
   req, err := http.NewRequest("GET", strings.TrimRight(cfg.Endpoint, "/")+"/models", nil)
   if err != nil { return "", err }
   req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)")