fixing streaming

packages/tracker/meshnet_tracker/server.py

@@ -1193,20 +1193,22 @@ def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int,
     Prefers the response usage block; falls back to content estimates.
     Completion stays capped by the request's max-tokens bound, as before.
     """
-    usage = _usage_split(payload)
-    prompt = (usage or {}).get("prompt")
-    completion = (usage or {}).get("completion")
-    if prompt is None:
-        prompt = _estimate_prompt_tokens(request_body) or 0
-    if completion is None:
-        total = (usage or {}).get("total")
-        if total is not None:
-            completion = max(0, total - prompt)
-        else:
-            completion = _observed_non_stream_completion_tokens(payload)
-    limit = _requested_completion_token_limit(request_body)
-    if limit is not None:
-        completion = min(completion, limit)
+    usage = _usage_split(payload)
+    prompt_estimate = _estimate_prompt_tokens(request_body) or 0
+    prompt = (usage or {}).get("prompt")
+    completion = (usage or {}).get("completion")
+    if prompt is None:
+        prompt = prompt_estimate
+    if completion is None:
+        total = (usage or {}).get("total")
+        if total is not None:
+            completion = max(0, total - prompt)
+        else:
+            completion = _observed_non_stream_completion_tokens(payload)
+    limit = _requested_completion_token_limit(request_body)
+    if limit is not None and completion > limit:
+        completion = min(completion, limit)
+        prompt = max(prompt, prompt_estimate)
     return max(0, prompt), max(0, completion)
 
 
@@ -1776,6 +1778,7 @@ def _tracker_log_proxy_progress(
     relay: bool = False,
 ) -> None:
     elapsed = time.monotonic() - started
+    effective_elapsed = max(elapsed, 1e-6)
     _tracker_log(
         server,
         "info",
@@ -1787,7 +1790,7 @@ def _tracker_log_proxy_progress(
         relay=relay or None,
         tokens=tokens,
         elapsed_seconds=round(elapsed, 4),
-        tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0,
+        tokens_per_sec=round(tokens / effective_elapsed, 4) if tokens > 0 else 0.0,
         route=_node_route_summary(route_nodes),
     )
 
@@ -2608,6 +2611,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     first, frames, started,
                     model, route_model, route_nodes, api_key, node_work,
                     request_body=body,
+                    request_id=request_id,
                 )
                 finish_proxy_inflight()
                 return
@@ -2802,7 +2806,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 self.end_headers()
                 try:
                     self.wfile.write(resp_body)
-                except BrokenPipeError:
+                except (BrokenPipeError, ConnectionResetError):
                     pass
         finish_proxy_inflight()
 
@@ -2819,11 +2823,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         The tracker sees end-to-end request duration, not per-hop timings, so
         each hop gets the same route-level observation for now. Per-hop telemetry
         can refine this later without changing the external stats shape.
-        """
-        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
-        if server.stats is None or total_tokens <= 0 or elapsed_seconds <= 0:
-            return
-        models = [m for m in (requested_model, route_model) if m]
+        """
+        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
+        if server.stats is None or total_tokens <= 0:
+            return
+        elapsed_seconds = max(elapsed_seconds, 1e-6)
+        models = [m for m in (requested_model, route_model) if m]
         if len(models) == 2 and models[0] == models[1]:
             models = [models[0]]
         for node in route_nodes:
@@ -2930,19 +2935,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         model: str,
         route_model: str,
         route_nodes: list,
-        api_key: str | None,
-        node_work: list,
-        request_body: dict,
-    ) -> None:
+        api_key: str | None,
+        node_work: list,
+        request_body: dict,
+        request_id: str,
+    ) -> None:
         """Forward a streamed relay response (US-036) to the client as SSE,
         billing with the same accounting as the direct stream path."""
         headers = first.get("headers") if isinstance(first.get("headers"), dict) else {}
         self.send_response(int(first.get("status", 200)))
         self.send_header("Content-Type", headers.get("Content-Type", "text/event-stream; charset=utf-8"))
-        self.send_header("Cache-Control", "no-cache")
-        self.end_headers()
-        stream_usage: dict | None = None
-        observed_stream_tokens = 0
+        self.send_header("Cache-Control", "no-cache")
+        self.end_headers()
+        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
+        stream_usage: dict | None = None
+        observed_stream_tokens = 0
         client_gone = False
         for frame in itertools.chain([first], frames):
             chunk = frame.get("chunk") or ""
@@ -2956,11 +2963,22 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                 except BrokenPipeError:
                     # Keep draining frames — the nodes did the work; bill it.
                     client_gone = True
-            for line in data.splitlines():
-                observed, usage = _stream_line_tokens(line)
-                observed_stream_tokens += observed
-                if usage is not None:
-                    stream_usage = usage
+            for line in data.splitlines():
+                observed, usage = _stream_line_tokens(line)
+                observed_stream_tokens += observed
+                if observed:
+                    _tracker_log_proxy_progress(
+                        server,
+                        request_id=request_id,
+                        model=model,
+                        route_model=route_model,
+                        tokens=observed_stream_tokens,
+                        started=started,
+                        route_nodes=route_nodes,
+                        relay=True,
+                    )
+                if usage is not None:
+                    stream_usage = usage
         elapsed = time.monotonic() - started
         in_tokens, out_tokens = _stream_billable_split(
             observed_stream_tokens, stream_usage, request_body
@@ -2969,12 +2987,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             model, route_model, in_tokens + out_tokens, elapsed, route_nodes
         )
         tokens = in_tokens + out_tokens
-        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
-        _tracker_log(
+        _tracker_log(
             server,
             "info",
             "proxy complete",
-            request_id=session_id,
+            request_id=request_id,
             model=model,
             route_model=route_model,
             status=200,