From e410b92e341fe3b40edc4246729ad78c461a54e4 Mon Sep 17 00:00:00 2001
From: Toby Chui <tobychui@users.noreply.github.com>
Date: Sat, 13 Jul 2024 23:12:23 +0800
Subject: [PATCH] Fixed #235

- Added flush sniffing for keep-alive request
- Set default flush interval to 100ms for hostname and 500ms for virtual directory
---
 src/mod/dynamicproxy/dpcore/flush.go         | 7 +++++++
 src/mod/dynamicproxy/loadbalance/upstream.go | 2 ++
 src/mod/dynamicproxy/router.go               | 2 ++
 src/web/components/httprp.html               | 1 -
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/mod/dynamicproxy/dpcore/flush.go b/src/mod/dynamicproxy/dpcore/flush.go
index d62e058..5758081 100644
--- a/src/mod/dynamicproxy/dpcore/flush.go
+++ b/src/mod/dynamicproxy/dpcore/flush.go
@@ -3,6 +3,7 @@ package dpcore
 import (
 	"mime"
 	"net/http"
+	"strings"
 	"time"
 )
 
@@ -17,6 +18,12 @@ func (p *ReverseProxy) getFlushInterval(req *http.Request, res *http.Response) t
 		return -1
 	}
 
+	// Fixed issue #235: Added auto detection for ollama / llm output stream
+	connectionHeader := req.Header["Connection"]
+	if len(connectionHeader) > 0 && strings.Contains(strings.Join(connectionHeader, ","), "keep-alive") {
+		return -1
+	}
+
 	//Cannot sniff anything. Use default value
 	return p.FlushInterval
 
diff --git a/src/mod/dynamicproxy/loadbalance/upstream.go b/src/mod/dynamicproxy/loadbalance/upstream.go
index 181e218..e04b7f4 100644
--- a/src/mod/dynamicproxy/loadbalance/upstream.go
+++ b/src/mod/dynamicproxy/loadbalance/upstream.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 
 	"imuslab.com/zoraxy/mod/dynamicproxy/dpcore"
 )
@@ -39,6 +40,7 @@ func (u *Upstream) StartProxy() error {
 
 	proxy := dpcore.NewDynamicProxyCore(path, "", &dpcore.DpcoreOptions{
 		IgnoreTLSVerification: u.SkipCertValidations,
+		FlushInterval:         100 * time.Millisecond,
 	})
 
 	u.proxy = proxy
diff --git a/src/mod/dynamicproxy/router.go b/src/mod/dynamicproxy/router.go
index 2f3c03e..07b716e 100644
--- a/src/mod/dynamicproxy/router.go
+++ b/src/mod/dynamicproxy/router.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"net/url"
 	"strings"
+	"time"
 
 	"imuslab.com/zoraxy/mod/dynamicproxy/dpcore"
 )
@@ -58,6 +59,7 @@ func (router *Router) PrepareProxyRoute(endpoint *ProxyEndpoint) (*ProxyEndpoint
 
 		proxy := dpcore.NewDynamicProxyCore(path, vdir.MatchingPath, &dpcore.DpcoreOptions{
 			IgnoreTLSVerification: vdir.SkipCertValidations,
+			FlushInterval:         500 * time.Millisecond,
 		})
 		vdir.proxy = proxy
 		vdir.parent = endpoint
diff --git a/src/web/components/httprp.html b/src/web/components/httprp.html
index b0114b2..ba5466f 100644
--- a/src/web/components/httprp.html
+++ b/src/web/components/httprp.html
@@ -234,7 +234,6 @@
         var payload = $(row).attr("payload");
         payload = JSON.parse(decodeURIComponent(payload));
         console.log(payload);
-        //console.log(payload);
         columns.each(function(index) {
             var column = $(this);
             var oldValue = column.text().trim();