fix(aws): cancel Bedrock requests on client disconnect

[Calcium-Ion]

⚠️ 提交说明 / PR Notice

Important

• 请提供人工撰写的简洁摘要，避免直接粘贴未经整理的 AI 输出。

📝 变更描述 / Description

AWS Bedrock 请求原先使用独立的后台 Context，客户端断开后，上游生成仍会继续，最终可能按照完整 Token 用量结算。

本次修改让 Bedrock SDK 请求继承客户端请求 Context，同时保留 RelayTimeout：

• 客户端断开时取消流式、非流式及 Nova 请求。
• 流式请求监听 Context 取消并及时关闭 EventStream，不再读取取消后的最终 usage。
• 客户端取消导致的初始调用错误停止重试；普通上游错误和独立的 RelayTimeout 保持原有重试行为。
• 中断后的流仍执行 Anthropic finalizer，保留取消前的部分 usage 用于结算。
• 调整 usage_billing_path 的日志判定，使审计标签与实际采用的 BillingUsage 一致，不改变计费数值。

🚀 变更类型 / Type of change

• 🐛 Bug 修复 (Bug fix) - 请关联对应 Issue，避免将设计取舍、理解偏差或预期不一致直接归类为 bug
• ✨ 新功能 (New feature) - 重大特性建议先通过 Issue 沟通
• ⚡ 性能优化 / 重构 (Refactor)
• 📝 文档更新 (Documentation)

🔗 关联任务 / Related Issue

• 暂无对应 Issue
• Related: fix: 客户端断连时取消上游请求并停止重试 #4500
• Related: fix: avoid stale stream writes after client disconnect #5710

✅ 提交前检查项 / Checklist

• 人工确认: 我已亲自整理并撰写此描述，没有直接粘贴未经处理的 AI 输出。
• 非重复提交: 我已搜索现有的 Issues 与 PRs，确认不是重复提交。
• Bug fix 说明: 若此 PR 标记为 Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。
• 变更理解: 我已理解这些更改的工作原理及可能影响。
• 范围聚焦: 本 PR 未包含任何与当前任务无关的代码改动。
• 本地验证: 已在本地运行并通过测试或手动验证，维护者可以据此复核结果。
• 安全合规: 代码中无敏感凭据，且符合项目代码规范。

📸 运行证明 / Proof of Work

验证范围包括 Context 继承、客户端取消、重试控制、流关闭、部分 usage 保留及日志路径判定。

$ go test ./relay/channel/aws ./relay/channel/claude ./service
ok  github.com/QuantumNous/new-api/relay/channel/aws
ok  github.com/QuantumNous/new-api/relay/channel/claude
ok  github.com/QuantumNous/new-api/service

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->

## Summary by CodeRabbit

* **Bug Fixes**
  * AWS requests now honor cancellation and configured timeouts across standard and streaming responses.
  * Canceled or timed-out requests are correctly treated as non-retryable.
  * Streaming responses stop cleanly on cancellation while preserving partial usage and billing details.
  * Billing paths now more accurately distinguish local, upstream, and estimated usage across supported providers.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

[cursor]

Bugbot is not enabled for your account, so this pull request was not reviewed.

Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs.

[coderabbitai]

Walkthrough

Changes

AWS invocation lifecycle

Layer / File(s)	Summary
AWS context and error propagation relay/channel/aws/relay-aws.go, relay/channel/aws/relay_aws_test.go	AWS handlers derive contexts from incoming requests and use shared error classification. Tests cover cancellation, retry behavior, and SDK request termination.
AWS stream cancellation and usage relay/channel/aws/relay-aws.go, relay/channel/aws/relay_aws_test.go	Streaming handlers stop on cancellation or channel closure, close the upstream stream, and preserve the specified partial or final usage behavior.

Billing path normalization

Layer / File(s)	Summary
Normalized billing path resolution service/billing_usage.go, service/text_quota_test.go	Billing path selection uses normalized usage data. Tests cover effective billing, estimated Gemini usage, and invalid Anthropic payload fallbacks.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant AWSHandler
  participant BedrockSDK
  participant EventStream

  Client->>AWSHandler: Send request with context
  AWSHandler->>BedrockSDK: Invoke using derived context
  BedrockSDK-->>AWSHandler: Return response or event stream
  EventStream-->>AWSHandler: Deliver chunks and usage
  Client-->>AWSHandler: Cancel request
  AWSHandler->>EventStream: Close upstream stream
  AWSHandler-->>Client: Stop output and finalize response

Loading

Suggested reviewers: yyhhyyyyyy

Poem

A rabbit hops through streams of light,
Cancels a call before the night.
Chunks are saved, the stream is closed,
Usage paths are neatly composed.
“Done,” says Bun, with ears held high.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the primary change: canceling AWS Bedrock requests when the client disconnects.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/aws-contex

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@relay/channel/aws/relay_aws_test.go`:
- Line 448: Update the final body-length assertion in the cancellation test to
verify that responseWriter.Body.Len() is at least bodyLengthBeforeCancel, rather
than requiring exact equality. Keep the existing partial-usage and "[DONE]"
assertions unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: abbeff12-d2e9-42bc-ae14-35b8e32597c6

📥 Commits

Reviewing files that changed from the base of the PR and between cfaba1d and f04db28.

📒 Files selected for processing (4)

• relay/channel/aws/relay-aws.go
• relay/channel/aws/relay_aws_test.go
• service/billing_usage.go
• service/text_quota_test.go

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Relax the exact body-length assertion.

bodyLengthBeforeCancel is captured inside Flush at the moment the body first contains partial. The producer has already written all three initial events to the pipe at that point. The handler can therefore write further SSE frames before it observes ctx.Done(), so the final body length is not deterministic. This assertion can fail intermittently.

The invariants that matter are already covered by the partial-usage assertions and by assert.NotContains(..., "[DONE]"). Assert a lower bound instead of exact equality.

💚 Proposed change

-	assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len())
+	assert.GreaterOrEqual(t, responseWriter.Body.Len(), bodyLengthBeforeCancel)

As per coding guidelines: "prefer deterministic table tests with explicit inputs and exact outputs, and avoid coverage-only, fake stress, timing, or implementation-detail tests".

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len())
	assert.GreaterOrEqual(t, responseWriter.Body.Len(), bodyLengthBeforeCancel)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@relay/channel/aws/relay_aws_test.go` at line 448, Update the final
body-length assertion in the cancellation test to verify that
responseWriter.Body.Len() is at least bodyLengthBeforeCancel, rather than
requiring exact equality. Keep the existing partial-usage and "[DONE]"
assertions unchanged.

Source: Coding guidelines