Fix socket leaks on shutdown, remove hardcoded fake thinking block, and add debug logs for requests and tools

gemini_web2api.py

@@ -62,6 +62,7 @@ DEFAULT_CONFIG = {
     # Request jitter (ms) - randomized delays to mimic human behavior
     "jitter_min_ms": 50,
     "jitter_max_ms": 300,
+    "debug_mode": False,
 }
 
 CONFIG = dict(DEFAULT_CONFIG)
@@ -403,6 +404,60 @@ class GeminiHTTPClient:
 
         return resp.read().decode("utf-8", errors="replace")
 
+    def post_stream(self, url: str, data: bytes, headers: dict, cookies: dict = None):
+        """POST request that yields streaming chunks."""
+        if self._session:
+            return self._post_curl_stream(url, data, headers, cookies)
+        else:
+            return self._post_urllib_stream(url, data, headers, cookies)
+
+    def _post_curl_stream(self, url: str, data: bytes, headers: dict, cookies: dict = None):
+        self._session.cookies.clear()
+        proxy = CONFIG.get("proxy")
+        proxies = {"http": proxy, "https": proxy} if proxy else None
+        
+        resp = self._session.post(
+            url,
+            data=data,
+            headers=dict(headers),
+            cookies=cookies or {},
+            proxies=proxies,
+            allow_redirects=True,
+            stream=True
+        )
+        if resp.status_code != 200:
+            raise Exception(f"HTTP {resp.status_code}")
+            
+        for line in resp.iter_lines():
+            if line:
+                yield line.decode("utf-8", errors="replace")
+
+    def _post_urllib_stream(self, url: str, data: bytes, headers: dict, cookies: dict = None):
+        all_headers = dict(headers)
+        if cookies:
+            cookie_str = "; ".join(f"{k}={v}" for k, v in cookies.items())
+            existing = all_headers.get("Cookie", "")
+            if existing:
+                all_headers["Cookie"] = existing + "; " + cookie_str
+            else:
+                all_headers["Cookie"] = cookie_str
+
+        req = urllib.request.Request(url, data=data, headers=all_headers, method="POST")
+        ctx = ssl.create_default_context()
+        proxy = CONFIG.get("proxy")
+        if proxy:
+            opener = urllib.request.build_opener(
+                urllib.request.ProxyHandler({"http": proxy, "https": proxy}),
+                urllib.request.HTTPSHandler(context=ctx)
+            )
+            resp = opener.open(req, timeout=CONFIG["request_timeout_sec"])
+        else:
+            resp = urllib.request.urlopen(req, context=ctx, timeout=CONFIG["request_timeout_sec"])
+            
+        for line in resp:
+            if line:
+                yield line.decode("utf-8", errors="replace")
+
     def close(self):
         if self._session:
             try:
@@ -424,7 +479,7 @@ def get_http_client() -> GeminiHTTPClient:
 
 # ─── Gemini Protocol ─────────────────────────────────────────────────────────
 
-def gemini_stream_generate(prompt: str, model_info: dict) -> str:
+def gemini_stream_generate(prompt: str, model_info: dict, stream: bool = False):
     """Send prompt to Gemini StreamGenerate with retry.
 
     Uses the x-goog-ext-525001261-jspb header for model selection
@@ -497,7 +552,10 @@ def gemini_stream_generate(prompt: str, model_info: dict) -> str:
                 time.sleep(CONFIG["retry_delay_sec"])
             apply_jitter()
 
-            return client.post(url, data=body, headers=headers, cookies=cookies)
+            if stream:
+                return client.post_stream(url, data=body, headers=headers, cookies=cookies)
+            else:
+                return client.post(url, data=body, headers=headers, cookies=cookies)
         except Exception as e:
             last_err = e
             if attempt < CONFIG["retry_attempts"] - 1:
@@ -509,8 +567,8 @@ def clean_gemini_text(text: str) -> str:
     """Remove internal code execution artifacts and image placeholders."""
     # Convert code execution blocks to standard markdown
     text = re.sub(
-        r'```(python|javascript|text)\?code_(?:reference|stdout)&code_event_index=\d+',
-        r'```\1', text
+        r'\?code_(?:reference|stdout)&code_event_index=\d+',
+        '', text
     )
     # Remove googleusercontent placeholder URLs (image gen/retrieval/collection)
     text = re.sub(
@@ -571,47 +629,64 @@ def _scan_complete_wrb_frames(buf: str) -> list:
     return frames
 
 
-def extract_response_text(raw: str, model_info: dict = None) -> str:
-    """Parse StreamGenerate response to extract final text.
-
-    Uses the upstream's bracket-depth frame scanner for robustness.
-    """
-    frames = _scan_complete_wrb_frames(raw)
+def gemini_stream_parse(stream_generator, model_info: dict = None):
+    """Consume network chunks, parse wrb.fr frames, and yield text deltas incrementally."""
+    buf = ""
+    emitted_raw = ""
+    first_chunk = True
+    
+    for chunk in stream_generator:
+        if not chunk: continue
+        buf += chunk
+        
+        frames = _scan_complete_wrb_frames(buf)
+        if not frames: continue
+        
+        # In stream context, the parser needs to extract the latest text
+        texts = []
+        for elem in frames:
+            try:
+                if not isinstance(elem, list) or len(elem) < 3 or elem[0] != "wrb.fr":
+                    continue
+                rp = elem[2]
+                if not isinstance(rp, str) or len(rp) < 50:
+                    continue
+                payload = json.loads(rp)
+                
+                if isinstance(payload, list) and len(payload) > 4 and payload[4]:
+                    for part in payload[4]:
+                        if isinstance(part, list) and len(part) > 1 and part[1]:
+                            if isinstance(part[1], list):
+                                for t in part[1]:
+                                    if isinstance(t, str) and len(t) > 0:
+                                        texts.append(t)
+            except (json.JSONDecodeError, IndexError, TypeError):
+                pass
 
-    texts = []
-    for elem in frames:
-        try:
-            if not isinstance(elem, list) or len(elem) < 3 or elem[0] != "wrb.fr":
-                continue
-            rp = elem[2]
-            if not isinstance(rp, str) or len(rp) < 50:
-                continue
-            payload = json.loads(rp)
+        current_full_text = ""
+        for t in reversed(texts):
+            if t.strip():
+                current_full_text = t
+                break
+                
+        if current_full_text == emitted_raw:
+            continue
             
-            # Extract final text
-            if isinstance(payload, list) and len(payload) > 4 and payload[4]:
-                for part in payload[4]:
-                    if isinstance(part, list) and len(part) > 1 and part[1]:
-                        if isinstance(part[1], list):
-                            for t in part[1]:
-                                if isinstance(t, str) and len(t) > 0:
-                                    texts.append(t)
-        except (json.JSONDecodeError, IndexError, TypeError):
-            pass
-
-    # Take the last non-empty text (final/most complete response)
-    text = ""
-    for t in reversed(texts):
-        if t.strip():
-            text = t
-            break
+        if current_full_text.startswith(emitted_raw):
+            raw_delta = current_full_text[len(emitted_raw):]
+            emitted_raw = current_full_text
             
-    final_response = clean_gemini_text(text)
-    
-    if model_info and model_info.get("think") == 0:
-        final_response = f"<think>\nthink_mode: 0\nwe do not really recieve thinking block from gemini.\n</think>\n\n{final_response}"
-        
-    return final_response
+            if raw_delta:
+                cleaned_delta = clean_gemini_text(raw_delta)
+                first_chunk = False
+                
+                if cleaned_delta:
+                    yield cleaned_delta
+
+def extract_response_text(raw: str, model_info: dict = None) -> str:
+    """Parse StreamGenerate response to extract final text. (Backwards compatible)"""
+    gen = gemini_stream_parse([raw], model_info)
+    return "".join(list(gen))
 
 
 # ─── OpenAI Format Helpers ───────────────────────────────────────────────────
@@ -794,17 +869,22 @@ class GeminiHandler(BaseHTTPRequestHandler):
             return None, None, err
         return pub_name, model_info, None
 
-    def _call_gemini(self, prompt, model_info, tools):
-        raw = gemini_stream_generate(prompt, model_info)
-        text = extract_response_text(raw, model_info)
-        tool_calls = None
-        if tools and text:
-            text, tool_calls = parse_tool_calls(text)
-        return text or "", tool_calls
+    def _call_gemini(self, prompt, model_info, tools, stream=False):
+        raw = gemini_stream_generate(prompt, model_info, stream=stream)
+        if stream:
+            return gemini_stream_parse(raw, model_info)
+        else:
+            text = extract_response_text(raw, model_info)
+            tool_calls = None
+            if tools and text:
+                text, tool_calls = parse_tool_calls(text)
+            return text or "", tool_calls
 
     def handle_chat(self, body: bytes):
         try:
             req = json.loads(body)
+            if CONFIG.get("debug_mode"):
+                log(f"DEBUG [CHAT] REQUEST: {json.dumps(req, ensure_ascii=False)[:2000]}")
         except json.JSONDecodeError as e:
             self.send_json({"error": {"message": f"Invalid JSON payload: {e}. Body received: {body.decode('utf-8', errors='replace')}"}}, 400)
             return
@@ -815,47 +895,102 @@ class GeminiHandler(BaseHTTPRequestHandler):
             return
 
         tools = req.get("tools")
+        if CONFIG.get("debug_mode"):
+            think_status = "Enabled" if model_info.get("think") == 0 else "Disabled"
+            log(f"DEBUG [CHAT] MODEL: {model_name} (Think Mode: {think_status})")
+            if tools:
+                log(f"DEBUG [CHAT] TOOLS PROVIDED: {len(tools)} tools")
+                
         prompt = messages_to_prompt(req.get("messages", []), tools)
         if not prompt.strip():
             self.send_json({"error": {"message": "empty prompt"}}, 400)
             return
 
+        is_stream = bool(req.get("stream"))
+        
         try:
-            text, tool_calls = self._call_gemini(prompt, model_info, tools)
+            # If tools are provided, we must collect full text first to parse them, so disable network streaming
+            if tools:
+                text, tool_calls = self._call_gemini(prompt, model_info, tools, stream=False)
+                if CONFIG.get("debug_mode"):
+                    log(f"DEBUG [CHAT] RESPONSE TEXT: {text}")
+                    log(f"DEBUG [CHAT] RESPONSE TOOLS: {tool_calls}")
+            else:
+                result = self._call_gemini(prompt, model_info, tools, stream=is_stream)
+                if not is_stream and CONFIG.get("debug_mode"):
+                    log(f"DEBUG [CHAT] RESPONSE: {result}")
         except Exception as e:
             self.send_json({"error": {"message": f"upstream error: {e}"}}, 502)
             return
 
         cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
-        msg = {"role": "assistant", "content": text or None}
-        if tool_calls:
-            msg["tool_calls"] = tool_calls
-        finish = "tool_calls" if tool_calls else "stop"
-
-        if req.get("stream"):
+        
+        if is_stream:
             self.send_response(200)
             self.send_header("Content-Type", "text/event-stream")
             self.send_header("Cache-Control", "no-cache")
             self.send_header("Access-Control-Allow-Origin", "*")
             self.end_headers()
-            chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
-                     "model": model_name, "choices": [{"index": 0, "delta": msg, "finish_reason": finish}]}
-            self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
+            
+            if tools:
+                # Tools were present, so we ran synchronously. We yield the tool calls in streaming format.
+                if tool_calls:
+                    for tc in tool_calls:
+                        chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
+                                 "model": model_name, "choices": [{"index": 0, "delta": {"tool_calls": [tc]}}]}
+                        self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
+                    chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
+                             "model": model_name, "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]}
+                    self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
+                else:
+                    msg = {"role": "assistant", "content": text or ""}
+                    chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
+                             "model": model_name, "choices": [{"index": 0, "delta": msg, "finish_reason": "stop"}]}
+                    self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
+            else:
+                # Real streaming
+                first = True
+                for delta in result:
+                    if CONFIG.get("debug_mode") and delta:
+                        log(f"DEBUG [CHAT] CHUNK: {delta}")
+                    msg = {"role": "assistant"} if first else {}
+                    if delta: msg["content"] = delta
+                    first = False
+                    chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
+                             "model": model_name, "choices": [{"index": 0, "delta": msg, "finish_reason": None}]}
+                    self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
+                    self.wfile.flush()
+                
+                chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()),
+                         "model": model_name, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
+                self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
+            
             self.wfile.write(b"data: [DONE]\n\n")
             self.wfile.flush()
         else:
+            if tools:
+                msg = {"role": "assistant", "content": text or None}
+                if tool_calls:
+                    msg["tool_calls"] = tool_calls
+                finish = "tool_calls" if tool_calls else "stop"
+            else:
+                msg = {"role": "assistant", "content": result or None}
+                finish = "stop"
+                
             self.send_json({
                 "id": cid, "object": "chat.completion", "created": int(time.time()),
                 "model": model_name,
                 "choices": [{"index": 0, "message": msg, "finish_reason": finish}],
-                "usage": {"prompt_tokens": len(prompt)//4, "completion_tokens": len(text)//4,
-                          "total_tokens": (len(prompt)+len(text))//4},
+                "usage": {"prompt_tokens": len(prompt)//4, "completion_tokens": 0,
+                          "total_tokens": len(prompt)//4},
             })
 
     def handle_responses(self, body: bytes):
         """OpenAI Responses API for Codex CLI compatibility."""
         try:
             req = json.loads(body)
+            if CONFIG.get("debug_mode"):
+                log(f"DEBUG [RESP] REQUEST: {json.dumps(req, ensure_ascii=False)[:2000]}")
         except json.JSONDecodeError as e:
             self.send_json({"error": {"message": f"Invalid JSON payload: {e}. Body received: {body.decode('utf-8', errors='replace')}"}}, 400)
             return
@@ -867,6 +1002,12 @@ class GeminiHandler(BaseHTTPRequestHandler):
 
         input_items = req.get("input", [])
         tools = req.get("tools")
+        
+        if CONFIG.get("debug_mode"):
+            think_status = "Enabled" if model_info.get("think") == 0 else "Disabled"
+            log(f"DEBUG [RESP] MODEL: {model_name} (Think Mode: {think_status})")
+            if tools:
+                log(f"DEBUG [RESP] TOOLS PROVIDED: {len(tools)} tools")
 
         messages = []
         if req.get("instructions"):
@@ -915,6 +1056,9 @@ class GeminiHandler(BaseHTTPRequestHandler):
 
         try:
             text, tool_calls = self._call_gemini(prompt, model_info, tools)
+            if CONFIG.get("debug_mode"):
+                log(f"DEBUG [RESP] RESPONSE TEXT: {text}")
+                log(f"DEBUG [RESP] RESPONSE TOOLS: {tool_calls}")
         except Exception as e:
             self.send_json({"error": {"message": f"upstream error: {e}"}}, 502)
             return
@@ -993,6 +1137,7 @@ def main():
     parser.add_argument("--config", type=str, default=None)
     parser.add_argument("--cookie-file", type=str, default=None, help="Path to cookie file")
     parser.add_argument("--proxy", type=str, default=None, help="HTTP proxy, e.g. http://127.0.0.1:7890")
+    parser.add_argument("--debug", action="store_true", help="Enable debug logging of requests/responses")
     parser.add_argument("--version", action="version", version=f"gemini-web2api {__version__}")
     args = parser.parse_args()
 
@@ -1010,6 +1155,8 @@ def main():
         CONFIG["cookie_file"] = args.cookie_file
     if args.proxy:
         CONFIG["proxy"] = args.proxy
+    if args.debug:
+        CONFIG["debug_mode"] = True
 
     # Initialize HTTP client
     get_http_client()
@@ -1032,13 +1179,14 @@ def main():
     print(f"  Proxy:     {CONFIG.get('proxy') or 'none (uses system env HTTP_PROXY/HTTPS_PROXY)'}")
     print(f"  Retry:     {CONFIG['retry_attempts']}x / {CONFIG['retry_delay_sec']}s")
     print(f"  Jitter:    {CONFIG['jitter_min_ms']}-{CONFIG['jitter_max_ms']}ms")
+    print(f"  Debug:     {'enabled' if CONFIG.get('debug_mode') else 'disabled'}")
     print()
     try:
         server.serve_forever()
     except KeyboardInterrupt:
         print("\nStopped.")
         get_http_client().close()
-        server.shutdown()
+        server.server_close()
 
 
 if __name__ == "__main__":