starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): record image usage and endpoint request models

Browse Source 
Persist models declared by standalone image and search requests without
routing them through the text model catalog. Extract image token usage
before delivering the unchanged JSON response, preserving unknown counts
and upstream totals without double-counting image token details.

Validation: 1,474 related backend tests, workspace rustfmt and strict
Clippy passed. One authorized live gpt-image-2 generation reported 18 input,
229 output and 247 total tokens, matching the canonical usage exactly.
Search responses inspected in official source and production logs did not
report usage, so their token counts remain unknown.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-08 02:23:00 +08:00
co-authored by Codex
parent f465a337e8
commit 91f35d88cf
8 changed files with 293 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/engine/execution.rs
+1 -1
View File
@@ -524,7 +524,7 @@ impl DefaultExecutionService {
operation: operation.kind(),
endpoint: metadata.endpoint,
client_transport: metadata.transport.as_str().to_owned(),
requested_model: target.into_public_model(),
requested_model: target.into_public_model().or(observation.requested_model),
client_ip: metadata.client_ip,
user_agent: metadata.user_agent,
reasoning_effort: observation.reasoning_effort,
backend/crates/gateway-core/src/engine/provider.rs
+5 -1
View File
@@ -20,7 +20,9 @@ use crate::event::{EventSequenceValidator, ProviderEvent};
use crate::identity::ProviderKind;
use crate::operation::Operation;
use crate::policy::ClientApiKeyId;
use crate::routing::{ModelCapabilities, ModelPresentation, ProviderCandidate, UpstreamModelId};
use crate::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, PublicModelId, UpstreamModelId,
};
use crate::upstream::OpaqueUpstreamValue;
use crate::upstream::{UpstreamSendState, UpstreamTransport};
@@ -416,6 +418,8 @@ pub struct ProviderRequest {
/// Provider 对公共观测表可解释的请求语义；未知字段保持空值。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProviderRequestObservation {
/// 独立 Provider 端点声明的请求模型，仅用于观测，不参与文本模型目录路由。
pub requested_model: Option<PublicModelId>,
/// 客户端原始请求中的推理强度。
pub reasoning_effort: Option<String>,
pub reasoning_preset: Option<String>,
backend/crates/gateway-core/tests/engine/execution.rs
+21 -5
View File
@@ -26,7 +26,7 @@ use gateway_core::engine::execution::{
use gateway_core::engine::probe::{AccountProbe, AccountProbeErrorSource, AccountProbeRequest};
use gateway_core::engine::provider::{
Provider, ProviderCallMetadata, ProviderCatalogGeneration, ProviderModelCapabilities,
ProviderRegistry, ProviderRequest, ProviderStream,
ProviderRegistry, ProviderRequest, ProviderRequestObservation, ProviderStream,
};
use gateway_core::engine::{
AttemptContext, AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization,
@@ -260,7 +260,9 @@ impl Provider for LocalFailingProvider {
}
}
struct ColdFailingProvider;
struct ColdFailingProvider {
requested_model: Option<PublicModelId>,
}
#[async_trait]
impl Provider for ColdFailingProvider {
@@ -272,6 +274,13 @@ impl Provider for ColdFailingProvider {
ProviderCatalogGeneration::default()
}
fn request_observation(&self, _: &Operation, _: &ClientApiKeyId) -> ProviderRequestObservation {
ProviderRequestObservation {
requested_model: self.requested_model.clone(),
..Default::default()
}
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
@@ -323,12 +332,19 @@ fn successful_authentication_should_record_client_key_usage() {
#[test]
fn provider_endpoint_should_persist_its_real_v1_endpoint() {
assert_provider_endpoint_observation(None);
assert_provider_endpoint_observation(Some("gpt-image-2"));
}
fn assert_provider_endpoint_observation(model: Option<&str>) {
let store = Arc::new(TrackingExecutionStore::default());
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(client_snapshot()),
store.clone(),
ProviderRegistry::new([Arc::new(ColdFailingProvider) as Arc<dyn Provider>])
.expect("provider registry"),
ProviderRegistry::new([Arc::new(ColdFailingProvider {
requested_model: model.map(|model| PublicModelId::new(model).expect("model")),
}) as Arc<dyn Provider>])
.expect("provider registry"),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
@@ -369,7 +385,7 @@ fn provider_endpoint_should_persist_its_real_v1_endpoint() {
vec!["/v1/images/generations".to_owned()],
"execution error: {error:?}"
);
assert_eq!(store.requested_models(), vec![None]);
assert_eq!(store.requested_models(), vec![model.map(str::to_owned)]);
let upstream_models = store.upstream_models();
assert!(!upstream_models.is_empty());
assert!(upstream_models.iter().all(Option::is_none));
backend/crates/providers/openai/src/provider/execution.rs
+34
View File
@@ -1,5 +1,7 @@
//! OpenAI attempt 的选择、发送与响应流执行。
use gateway_core::metering::Usage;
use super::*;
impl CodexProvider {
@@ -490,6 +492,11 @@ pub(super) fn cold_json_response_stream(request: ColdJsonResponse) -> EventStrea
let response_meta =
ResponseMeta::for_provider_endpoint(request.context.request_id().as_str());
yield ProviderEvent::canonical(GatewayEvent::Started(response_meta.clone()));
if matches!(request.endpoint_path, CODEX_IMAGE_GENERATIONS_PATH | CODEX_IMAGE_EDITS_PATH)
&& let Some(usage) = image_response_usage(&response.body)
{
yield ProviderEvent::canonical(GatewayEvent::Usage(usage));
}
let wire = ProtocolWireEvent::raw_json(PROVIDER_NAME, response.body).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)
})?;
@@ -500,6 +507,33 @@ pub(super) fn cold_json_response_stream(request: ColdJsonResponse) -> EventStrea
})
}
fn image_response_usage(body: &[u8]) -> Option<Usage> {
// 只保留 usage，跳过通常很大的 base64 图片；原始响应仍按字节透传。
#[derive(Deserialize)]
struct ImageUsageEnvelope {
usage: Option<Value>,
}
let raw = serde_json::from_slice::<ImageUsageEnvelope>(body)
.ok()?
.usage?;
let mut usage = Usage::new();
usage.input_tokens = raw.get("input_tokens").and_then(Value::as_u64);
usage.output_tokens = raw.get("output_tokens").and_then(Value::as_u64);
usage.cached_tokens = raw
.pointer("/input_tokens_details/cached_tokens")
.and_then(Value::as_u64);
usage.image_input_tokens = raw
.pointer("/input_tokens_details/image_tokens")
.and_then(Value::as_u64);
usage.image_output_tokens = raw
.pointer("/output_tokens_details/image_tokens")
.and_then(Value::as_u64);
// 总量是上游独立报告的事实；图片明细是总输入/输出的子集，不能再次相加。
usage.total_tokens = raw.get("total_tokens").and_then(Value::as_u64);
(usage != Usage::default()).then_some(usage)
}
pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
let ColdResponse {
client,
backend/crates/providers/openai/src/provider/mod.rs
+10
View File
@@ -237,6 +237,15 @@ impl Provider for CodexProvider {
_ => None,
};
return ProviderRequestObservation {
requested_model: match operation {
Operation::Search(request) => {
observation::endpoint_requested_model(request.payload())
}
Operation::GenerateImage(request) => {
observation::endpoint_requested_model(request.payload())
}
_ => None,
},
continuation: ContinuationRequestObservation {
affinity_hash: affinity.map(|affinity| affinity.persistence_hash().to_owned()),
..Default::default()
@@ -259,6 +268,7 @@ impl Provider for CodexProvider {
requested: previous_response_id.is_some(),
};
ProviderRequestObservation {
requested_model: None,
reasoning_effort,
reasoning_preset: semantics.reasoning_preset.map(str::to_owned),
request_kind: semantics.request_kind,
backend/crates/providers/openai/src/provider/observation.rs
+17
View File
@@ -2,6 +2,23 @@
use super::*;
pub(super) fn endpoint_requested_model(
payload: &gateway_core::operation::RawJsonPayload,
) -> Option<gateway_core::routing::PublicModelId> {
if payload.protocol() != PROVIDER_NAME {
return None;
}
// 只读取 model，避免把编辑请求中的整张图片复制到观测数据。
#[derive(Deserialize)]
struct RequestModel {
model: Option<Value>,
}
let model = serde_json::from_slice::<RequestModel>(payload.body())
.ok()?
.model?;
gateway_core::routing::PublicModelId::new(model.as_str()?.to_owned()).ok()
}
/// OpenAI Responses 的观测状态完全归 Provider 所有。
///
/// 每次原始 SSE/WS 事件推进状态后重新生成不可变 observation，Core 只负责携带、
backend/crates/providers/openai/tests/provider/contract.rs
+204
View File
@@ -26,6 +26,7 @@ use gateway_core::error::{
};
use gateway_core::event::{GatewayEvent, WebSocketPoolKind};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::metering::Usage;
use gateway_core::operation::{
CapabilityRequirements, GenerateRequest, ImageRequest, ImageRequestKind, Operation,
OperationKind, ProtocolPayload, ProviderSessionState, RawJsonPayload, StandaloneSearchRequest,
@@ -1051,6 +1052,13 @@ async fn image_endpoints_bypass_only_the_text_catalog_and_preserve_the_current_c
let mut observed_http_json = false;
while let Some(event) = stream.next().await {
let event = event.expect("image provider event");
assert!(
event
.canonical_facts()
.iter()
.all(|fact| !matches!(fact, GatewayEvent::Usage(_))),
"missing upstream usage must remain unknown"
);
completed |= event
.canonical_facts()
.iter()
@@ -1075,6 +1083,156 @@ async fn image_endpoints_bypass_only_the_text_catalog_and_preserve_the_current_c
server.verify().await;
}
#[tokio::test]
async fn image_endpoints_should_report_usage_before_delivering_the_unchanged_body() {
let response = br#"{ "created":1778832973,"data":[{"b64_json":"AAEC"}],"quality":"medium","size":"1024x1536","usage":{"input_tokens":1474,"input_tokens_details":{"image_tokens":1457,"text_tokens":17,"cached_tokens":100},"output_tokens":1372,"output_tokens_details":{"image_tokens":1372,"text_tokens":0},"total_tokens":2846},"future":9007199254740993 }"#;
let expected = Usage {
input_tokens: Some(1474),
output_tokens: Some(1372),
cached_tokens: Some(100),
image_input_tokens: Some(1457),
image_output_tokens: Some(1372),
total_tokens: Some(2846),
..Usage::default()
};
for kind in [ImageRequestKind::Generation, ImageRequestKind::Edit] {
let events = image_usage_events(kind, response).await;
let usage = events
.iter()
.flat_map(|event| event.canonical_facts())
.filter_map(|fact| match fact {
GatewayEvent::Usage(usage) => Some(usage.clone()),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(usage.as_slice(), std::slice::from_ref(&expected));
let usage_index = events
.iter()
.position(|event| {
event
.canonical_facts()
.iter()
.any(|fact| matches!(fact, GatewayEvent::Usage(_)))
})
.expect("usage event");
let wire_index = events
.iter()
.position(|event| {
event
.wire_event()
.and_then(|wire| wire.raw_json_body())
.is_some()
})
.expect("raw image response");
assert!(
usage_index < wire_index,
"usage must be observed before image delivery"
);
assert_eq!(
events[wire_index]
.wire_event()
.and_then(|wire| wire.raw_json_body())
.map(|body| body.as_ref()),
Some(response.as_slice())
);
}
}
#[tokio::test]
async fn image_usage_should_preserve_unknown_fields_and_explicit_zero_counts() {
let cases = [
(json!(null), None),
(json!({}), None),
(json!([]), None),
(
json!({"input_tokens": -1, "output_tokens": "42", "total_tokens": 1.5}),
None,
),
(
json!({"input_tokens": 0, "output_tokens": 0, "total_tokens": 0, "input_tokens_details": {"image_tokens": 0, "cached_tokens": 0}, "output_tokens_details": {"image_tokens": 0}}),
Some(Usage {
input_tokens: Some(0),
output_tokens: Some(0),
total_tokens: Some(0),
cached_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
..Usage::default()
}),
),
(
json!({"input_tokens": 17, "output_tokens": 1372, "input_tokens_details": {"image_tokens": "invalid"}, "output_tokens_details": {"image_tokens": 1372}}),
Some(Usage {
input_tokens: Some(17),
output_tokens: Some(1372),
image_output_tokens: Some(1372),
..Usage::default()
}),
),
];
for (usage, expected) in cases {
let response = serde_json::to_vec(
&json!({"created":1778832973,"data":[{"b64_json":"AAEC"}],"usage":usage}),
)
.expect("image response JSON");
let events = image_usage_events(ImageRequestKind::Generation, &response).await;
let actual = events
.iter()
.flat_map(|event| event.canonical_facts())
.filter_map(|fact| match fact {
GatewayEvent::Usage(usage) => Some(usage.clone()),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(actual, expected.into_iter().collect::<Vec<_>>());
let raw = events
.iter()
.find_map(|event| event.wire_event().and_then(|wire| wire.raw_json_body()));
assert_eq!(raw.map(|body| body.as_ref()), Some(response.as_slice()));
}
}
async fn image_usage_events(
kind: ImageRequestKind,
response: &[u8],
) -> Vec<gateway_core::event::ProviderEvent> {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
let endpoint = match kind {
ImageRequestKind::Generation => "/codex/images/generations",
ImageRequestKind::Edit => "/codex/images/edits",
};
Mock::given(method("POST"))
.and(path(endpoint))
.respond_with(
ResponseTemplate::new(200).set_body_raw(response.to_vec(), "application/json"),
)
.expect(1)
.mount(&server)
.await;
let provider = provider_with_base_url(&store, server.uri());
let payload = RawJsonPayload::new(
"openai",
Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"draw a tree"}"#),
)
.expect("image payload");
let operation = Operation::GenerateImage(ImageRequest::from_raw_json(kind, payload));
let mut stream = provider
.execute(
planned_provider_endpoint_request("openai", operation),
context("req_image_usage", CancellationToken::new()),
)
.await
.expect("prepare image stream");
let mut events = Vec::new();
while let Some(event) = stream.next().await {
events.push(event.expect("image provider event"));
}
server.verify().await;
events
}
#[tokio::test]
async fn image_endpoint_returns_the_exact_upstream_error_response() {
let store = Arc::new(MemoryAccountStore::default());
@@ -5631,6 +5789,52 @@ fn request_observation_reads_openai_metadata_without_changing_the_operation() {
assert!(matches!(operation, Operation::Generate(_)));
}
#[test]
fn endpoint_observation_should_read_models_without_rewriting_or_requiring_a_catalog() {
let store = Arc::new(MemoryAccountStore::default());
let provider = provider(&store);
let client_key = ClientApiKeyId::new("key_endpoint_model").expect("client key");
for model in [
json!("gpt-image-future"),
json!("gpt-5.6-sol"),
json!(null),
json!(42),
json!(""),
] {
let body = serde_json::to_vec(&json!({"model":model,"future":9007199254740993_u64}))
.expect("body");
let payload =
RawJsonPayload::new("openai", Bytes::copy_from_slice(&body)).expect("payload");
let operations = [
Operation::GenerateImage(ImageRequest::from_raw_json(
ImageRequestKind::Generation,
payload.clone(),
)),
Operation::GenerateImage(ImageRequest::from_raw_json(
ImageRequestKind::Edit,
payload.clone(),
)),
Operation::Search(StandaloneSearchRequest::from_raw_json(payload)),
];
for operation in operations {
let observation = provider.request_observation(&operation, &client_key);
assert_eq!(
observation
.requested_model
.as_ref()
.map(PublicModelId::as_str),
model.as_str().filter(|model| !model.is_empty()),
);
let payload = match &operation {
Operation::GenerateImage(request) => request.payload(),
Operation::Search(request) => request.payload(),
_ => unreachable!(),
};
assert_eq!(payload.body().as_ref(), body.as_slice());
}
}
}
#[test]
fn request_observation_preserves_the_raw_reasoning_effort() {
let store = Arc::new(MemoryAccountStore::default());
backend/crates/providers/xai/src/provider/mod.rs
+1
View File
@@ -152,6 +152,7 @@ impl Provider for GrokBuildProvider {
}
let semantics = codex_responses_request_semantics(payload.body(), payload.context());
ProviderRequestObservation {
requested_model: None,
reasoning_effort: semantics.reasoning_effort,
reasoning_preset: semantics.reasoning_preset.map(str::to_owned),
request_kind: semantics.request_kind,