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 中重复的条件判断
- 新增对应单元测试
This commit is contained in:
fawney19
2026-03-17 03:37:06 +08:00
parent 40e0b82fa0
commit 438f16094f
6 changed files with 159 additions and 22 deletions

View File

@@ -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 "")