refactor(stream): 将不完整流 token 估算逻辑收敛到 StreamContext

- 新增 has_partial_response / ensure_estimated_output_tokens / should_estimate_incomplete_tokens 方法
- CancelledError 路径在归因前即补充 output_tokens，确保日志包含估算值
- CLI Handler 和 Chat Handler 的兜底估算统一使用 should_estimate_incomplete_tokens
- 移除 cli_monitor_mixin 和 stream_telemetry 中重复的条件判断
- 新增对应单元测试

src/api/handlers/base/cli_monitor_mixin.py

@@ -170,6 +170,8 @@ class CliMonitorMixin:
             # 也可能是服务端（重载/关停/内部取消）导致的协程取消。
             # 这里尽量做一次"断连归因"：仅当能确认客户端已断开时才记为 499 cancelled。
             time_since_last_chunk = time_module.time() - last_chunk_time
+            if not ctx.has_completion:
+                ctx.ensure_estimated_output_tokens()
 
             is_client_disconnected = False
             disconnect_check_uncertain = False
@@ -287,6 +289,11 @@ class CliMonitorMixin:
                     bg_db, user, api_key, ctx.request_id, self.client_ip
                 )
 
+                if ctx.should_estimate_incomplete_tokens():
+                    self._estimate_tokens_for_incomplete_stream(
+                        ctx, ctx.provider_request_body or original_request_body
+                    )
+
                 response_body = ctx.build_response_body(response_time_ms)
                 client_response_body = ctx.build_client_response_body(response_time_ms)
 
@@ -398,17 +405,6 @@ class CliMonitorMixin:
 
                     # 流未正常完成（如上游截断/连接中断）且无 token 数据时，
                     # 从已收集的文本和请求体估算 tokens，避免 usage 记录为 0
-                    if (
-                        not ctx.has_completion
-                        and ctx.data_count > 0
-                        and ctx.input_tokens == 0
-                        and ctx.output_tokens == 0
-                    ):
-                        # 用实际发给 Provider 的请求体估算 token（格式转换时与客户端请求体不同）
-                        self._estimate_tokens_for_incomplete_stream(
-                            ctx, ctx.provider_request_body or original_request_body
-                        )
-
                     # 流式成功时，返回给客户端的是提供商响应头 + SSE 必需头
                     client_response_headers = filter_proxy_response_headers(ctx.response_headers)
                     client_response_headers.update(

src/api/handlers/base/stream_context.py

@@ -331,6 +331,29 @@ class StreamContext:
         """检查是否因客户端断开连接而结束"""
         return self.status_code == 499
 
+    def has_partial_response(self) -> bool:
+        """是否已收到部分流式响应数据。"""
+        return self.data_count > 0 or self.chunk_count > 0 or self.collected_text_length > 0
+
+    def ensure_estimated_output_tokens(self) -> bool:
+        """在缺少 usage 时，基于已收集文本补充输出 tokens。"""
+        if self.output_tokens > 0 or self.collected_text_length <= 0:
+            return False
+        self.output_tokens = max(1, self.collected_text_length // 4)
+        return True
+
+    def should_estimate_incomplete_tokens(self) -> bool:
+        """流异常结束且尚无 usage 时，是否应做兜底 token 估算。
+
+        使用 or 而非 and：CancelledError 路径中 ensure_estimated_output_tokens
+        可能已补了 output_tokens，但 input_tokens 仍为 0，此时仍需估算。
+        """
+        return (
+            not self.has_completion
+            and (self.input_tokens == 0 or self.output_tokens == 0)
+            and self.has_partial_response()
+        )
+
     def set_ttfb_ms(self, ms: int) -> None:
         """将首字节响应耗时（TTFB）注入到 proxy_info 中"""
         if self.proxy_info is not None:

src/api/handlers/base/stream_telemetry.py

@@ -100,16 +100,9 @@ class StreamTelemetryRecorder:
                 writer = await self._get_telemetry_writer(bg_db, ctx, response_time_ms)
                 if writer is None:
                     return
-                # 兜底估算：流未正常完成且 token 均为 0 时，从请求体粗略估算
-                # 覆盖 Chat Handler 路径（CLI Handler 在更早的位置已做估算，
-                # 若已估算过则 token > 0，此处条件不会触发）
-                if (
-                    ctx.is_success()
-                    and not ctx.has_completion
-                    and ctx.data_count > 0
-                    and ctx.input_tokens == 0
-                    and ctx.output_tokens == 0
-                ):
+                # 兜底估算：流未正常完成且 token 均为 0 时，从请求体粗略估算。
+                # 覆盖成功但缺少 completion，以及已传出部分数据后被中断的场景。
+                if ctx.should_estimate_incomplete_tokens():
                     # 用实际发给 Provider 的请求体估算 token（格式转换时与客户端请求体不同）
                     self._estimate_tokens_for_incomplete_stream(
                         ctx, ctx.provider_request_body or original_request_body

tests/api/handlers/base/test_cli_monitor_mixin.py

@@ -33,6 +33,12 @@ async def _cancel_immediately() -> AsyncGenerator[bytes, None]:
     raise asyncio.CancelledError()
 
 
+async def _yield_once_then_cancel(ctx: StreamContext) -> AsyncGenerator[bytes, None]:
+    ctx.append_text("partial output")
+    yield b"data: chunk\n\n"
+    raise asyncio.CancelledError()
+
+
 @pytest.mark.asyncio
 async def test_create_monitored_stream_marks_client_disconnected_when_confirmed() -> None:
     monitor = _DummyMonitor()
@@ -70,7 +76,9 @@ async def test_create_monitored_stream_marks_server_cancelled_when_confirmed_con
 
 
 @pytest.mark.asyncio
-async def test_create_monitored_stream_marks_cancelled_unknown_when_disconnect_check_uncertain() -> None:
+async def test_create_monitored_stream_marks_cancelled_unknown_when_disconnect_check_uncertain() -> (
+    None
+):
     monitor = _DummyMonitor()
     monitor.CANCEL_DISCONNECT_RETRY_DELAYS_SECONDS = ()
     ctx = StreamContext(model="test-model", api_format="openai:cli", request_id="req-unknown")
@@ -85,3 +93,23 @@ async def test_create_monitored_stream_marks_cancelled_unknown_when_disconnect_c
     assert ctx.status_code == 503
     assert ctx.error_message == "cancelled_unknown"
     assert "cancel_origin=cancelled_unknown" in (ctx.upstream_response or "")
+
+
+@pytest.mark.asyncio
+async def test_create_monitored_stream_estimates_output_tokens_before_unknown_cancel_log() -> None:
+    monitor = _DummyMonitor()
+    monitor.CANCEL_DISCONNECT_RETRY_DELAYS_SECONDS = ()
+    ctx = StreamContext(model="test-model", api_format="openai:cli", request_id="req-estimate")
+
+    request = _RequestStub([asyncio.TimeoutError()])
+    monitored = monitor._create_monitored_stream(ctx, _yield_once_then_cancel(ctx), request)
+
+    with pytest.raises(asyncio.CancelledError):
+        async for _ in monitored:
+            pass
+
+    expected_output_tokens = max(1, len("partial output") // 4)
+    assert ctx.status_code == 503
+    assert ctx.error_message == "cancelled_unknown"
+    assert ctx.output_tokens == expected_output_tokens
+    assert f"output_tokens={expected_output_tokens}" in (ctx.upstream_response or "")

tests/api/handlers/base/test_stream_context.py

@@ -129,3 +129,34 @@ def test_get_log_summary_without_first_byte_time() -> None:
     assert "TTFB:" not in summary
     assert "Total: 456ms" in summary
     assert "in:100 out:50" in summary
+
+
+def test_ensure_estimated_output_tokens_uses_collected_text() -> None:
+    ctx = StreamContext(model="test-model", api_format="openai:chat")
+    ctx.append_text("partial output")
+
+    changed = ctx.ensure_estimated_output_tokens()
+
+    assert changed is True
+    assert ctx.output_tokens == max(1, len("partial output") // 4)
+
+
+def test_should_estimate_incomplete_tokens_for_interrupted_partial_stream() -> None:
+    ctx = StreamContext(model="test-model", api_format="openai:chat")
+    ctx.status_code = 503
+    ctx.chunk_count = 3
+
+    assert ctx.should_estimate_incomplete_tokens() is True
+
+
+def test_should_estimate_incomplete_tokens_when_output_already_estimated() -> None:
+    """ensure_estimated_output_tokens 已补了 output，但 input 仍为 0 时仍需估算。"""
+    ctx = StreamContext(model="test-model", api_format="openai:chat")
+    ctx.status_code = 503
+    ctx.chunk_count = 3
+    ctx.append_text("partial")
+    ctx.ensure_estimated_output_tokens()
+
+    assert ctx.output_tokens > 0
+    assert ctx.input_tokens == 0
+    assert ctx.should_estimate_incomplete_tokens() is True

tests/api/handlers/base/test_stream_telemetry.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from src.api.handlers.base import stream_telemetry as stream_telemetry_module
+from src.api.handlers.base.stream_context import StreamContext
+from src.api.handlers.base.stream_telemetry import StreamTelemetryRecorder
+
+
+class _DummyDb:
+    def close(self) -> None:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_record_stream_stats_estimates_tokens_for_failed_partial_stream(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    recorder = StreamTelemetryRecorder(
+        request_id="req-telemetry",
+        user_id="1",
+        api_key_id="2",
+        client_ip="127.0.0.1",
+        format_id="openai:chat",
+    )
+    recorder._get_telemetry_writer = AsyncMock(  # type: ignore[method-assign]
+        return_value=SimpleNamespace(include_bodies=False)
+    )
+    recorder._dispatch_record = AsyncMock()  # type: ignore[method-assign]
+    recorder._update_candidate_status = AsyncMock()  # type: ignore[method-assign]
+
+    ctx = StreamContext(
+        model="test-model",
+        api_format="openai:chat",
+        request_id="req-telemetry",
+        user_id=1,
+        api_key_id=2,
+    )
+    ctx.provider_name = "test-provider"
+    ctx.status_code = 503
+    ctx.data_count = 2
+    ctx.chunk_count = 4
+    ctx.append_text("partial output")
+
+    monkeypatch.setattr(stream_telemetry_module, "get_db", lambda: iter([_DummyDb()]))
+    monkeypatch.setattr(
+        stream_telemetry_module.SystemConfigService,
+        "should_log_body",
+        lambda _db: False,
+    )
+    monkeypatch.setattr(stream_telemetry_module.config, "stream_stats_delay", 0)
+
+    await recorder.record_stream_stats(
+        ctx,
+        original_headers={},
+        original_request_body={"input": [{"content": "hello world"}]},
+        start_time=time.time(),
+    )
+
+    assert ctx.input_tokens > 0
+    assert ctx.output_tokens == max(1, len("partial output") // 4)
+    recorder._dispatch_record.assert_awaited_once()  # type: ignore[attr-defined]