diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index dc3254a7..2832f41c 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -30,8 +30,9 @@ import ( ) const ( - codexUserAgent = "codex-tui/0.118.0 (Mac OS 26.3.1; arm64) iTerm.app/3.6.9 (codex-tui; 0.118.0)" - codexOriginator = "codex-tui" + codexUserAgent = "codex-tui/0.118.0 (Mac OS 26.3.1; arm64) iTerm.app/3.6.9 (codex-tui; 0.118.0)" + codexOriginator = "codex-tui" + codexDefaultImageToolModel = "gpt-image-2" ) var dataTag = []byte("data:") @@ -263,6 +264,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if detail, ok := helps.ParseCodexUsage(eventData); ok { reporter.Publish(ctx, detail) } + publishCodexImageToolUsage(ctx, reporter, body, eventData) completedData := eventData outputResult := gjson.GetBytes(completedData, "response.output") @@ -496,6 +498,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if detail, ok := helps.ParseCodexUsage(data); ok { reporter.Publish(ctx, detail) } + publishCodexImageToolUsage(ctx, reporter, body, data) data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback) translatedLine = append([]byte("data: "), data...) } @@ -859,6 +862,31 @@ func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth return body } +func publishCodexImageToolUsage(ctx context.Context, reporter *helps.UsageReporter, body []byte, completedData []byte) { + detail, ok := helps.ParseCodexImageToolUsage(completedData) + if !ok { + return + } + reporter.EnsurePublished(ctx) + reporter.PublishAdditionalModel(ctx, codexImageGenerationToolModel(body), detail) +} + +func codexImageGenerationToolModel(body []byte) string { + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + for _, tool := range tools.Array() { + if tool.Get("type").String() != "image_generation" { + continue + } + if model := strings.TrimSpace(tool.Get("model").String()); model != "" { + return model + } + break + } + } + return codexDefaultImageToolModel +} + func isCodexModelCapacityError(errorBody []byte) bool { if len(errorBody) == 0 { return false diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 97c1c611..615b6bed 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -48,6 +48,18 @@ func (r *UsageReporter) Publish(ctx context.Context, detail usage.Detail) { r.publishWithOutcome(ctx, detail, false) } +func (r *UsageReporter) PublishAdditionalModel(ctx context.Context, model string, detail usage.Detail) { + if r == nil { + return + } + model = strings.TrimSpace(model) + if model == "" { + return + } + detail = normalizeUsageDetailTotal(detail) + usage.PublishRecord(ctx, r.buildRecordForModel(model, detail, false)) +} + func (r *UsageReporter) PublishFailure(ctx context.Context) { r.publishWithOutcome(ctx, usage.Detail{}, true) } @@ -65,15 +77,20 @@ func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Det if r == nil { return } + detail = normalizeUsageDetailTotal(detail) + r.once.Do(func() { + usage.PublishRecord(ctx, r.buildRecord(detail, failed)) + }) +} + +func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail { if detail.TotalTokens == 0 { total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens if total > 0 { detail.TotalTokens = total } } - r.once.Do(func() { - usage.PublishRecord(ctx, r.buildRecord(detail, failed)) - }) + return detail } // ensurePublished guarantees that a usage record is emitted exactly once. @@ -93,9 +110,16 @@ func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool) usage.Reco if r == nil { return usage.Record{Detail: detail, Failed: failed} } + return r.buildRecordForModel(r.model, detail, failed) +} + +func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool) usage.Record { + if r == nil { + return usage.Record{Model: model, Detail: detail, Failed: failed} + } return usage.Record{ Provider: r.provider, - Model: r.model, + Model: model, Source: r.source, APIKey: r.apiKey, AuthID: r.authID, @@ -201,18 +225,15 @@ func ParseCodexUsage(data []byte) (usage.Detail, bool) { if !usageNode.Exists() { return usage.Detail{}, false } - detail := usage.Detail{ - InputTokens: usageNode.Get("input_tokens").Int(), - OutputTokens: usageNode.Get("output_tokens").Int(), - TotalTokens: usageNode.Get("total_tokens").Int(), + return parseOpenAIStyleUsageNode(usageNode), true +} + +func ParseCodexImageToolUsage(data []byte) (usage.Detail, bool) { + usageNode := gjson.ParseBytes(data).Get("response.tool_usage.image_gen") + if !usageNode.Exists() || !usageNode.IsObject() { + return usage.Detail{}, false } - if cached := usageNode.Get("input_tokens_details.cached_tokens"); cached.Exists() { - detail.CachedTokens = cached.Int() - } - if reasoning := usageNode.Get("output_tokens_details.reasoning_tokens"); reasoning.Exists() { - detail.ReasoningTokens = reasoning.Int() - } - return detail, true + return parseOpenAIStyleUsageNode(usageNode), true } func ParseOpenAIUsage(data []byte) usage.Detail { @@ -220,6 +241,10 @@ func ParseOpenAIUsage(data []byte) usage.Detail { if !usageNode.Exists() { return usage.Detail{} } + return parseOpenAIStyleUsageNode(usageNode) +} + +func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail { inputNode := usageNode.Get("prompt_tokens") if !inputNode.Exists() { inputNode = usageNode.Get("input_tokens")