release: prepare v0.1.137

This commit is contained in:
kone
2026-05-17 06:19:56 +08:00
parent f4055c773c
commit dd2b08d875
26 changed files with 344 additions and 32 deletions
+78 -8
View File
@@ -51,10 +51,11 @@ var (
)
type Usage struct {
InputTokens int
OutputTokens int
TotalTokens int
CacheReadInputTokens int
InputTokens int
OutputTokens int
TotalTokens int
CacheReadInputTokens int
CacheCreationInputTokens int
}
type StreamResult struct {
@@ -368,6 +369,8 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
inThinkingBlock := false
stripThinkingLeadingNewline := false
sawNonThinkingBlock := false
estimatedOutputTokens := 0
contextInputTokens := 0
writeEvent := func(event string, data any) error {
payload, err := json.Marshal(data)
@@ -488,6 +491,7 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
if toolUseID == "" || !streamingToolStarted[toolUseID] || streamingToolStopped[toolUseID] {
return nil
}
estimatedOutputTokens += countKiroTextTokens(fragment)
return writeEvent("content_block_delta", map[string]any{
"type": "content_block_delta",
"index": streamingToolBlockIndices[toolUseID],
@@ -571,6 +575,7 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
return err
}
}
estimatedOutputTokens += countKiroTextTokens(text)
return writeEvent("content_block_delta", map[string]any{
"type": "content_block_delta",
"index": contentBlockIndex,
@@ -611,6 +616,7 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
return err
}
inputJSON, _ := json.Marshal(tool.Input)
estimatedOutputTokens += estimateKiroOutputTokens("", []KiroToolUse{tool})
if err := writeEvent("content_block_delta", map[string]any{
"type": "content_block_delta",
"index": contentBlockIndex,
@@ -678,6 +684,7 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
return err
}
}
estimatedOutputTokens += countKiroTextTokens(text)
return writeEvent("content_block_delta", map[string]any{
"type": "content_block_delta",
"index": thinkingBlockIndex,
@@ -883,6 +890,14 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
if err := processStreamingToolUseEvent(event); err != nil {
return nil, err
}
case "contextUsageEvent":
contextUsage := nestedEvent(event, "contextUsageEvent")
if contextUsagePercent, ok := firstFloat(contextUsage, "contextUsagePercentage", "context_usage_percentage"); ok {
contextInputTokens = int(contextUsagePercent * float64(kiroContextWindowSize(model)) / 100)
if contextUsagePercent >= 100 && stopReason == "" {
stopReason = "model_context_window_exceeded"
}
}
case "messageMetadataEvent", "metadataEvent", "supplementaryWebLinksEvent", "usageEvent", "messageStopEvent", "message_stop":
updateUsageFromEvent(&usage, msg.EventType, event)
default:
@@ -912,6 +927,12 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
if err := closeThinking(); err != nil {
return nil, err
}
if usage.OutputTokens == 0 && estimatedOutputTokens > 0 {
usage.OutputTokens = estimatedOutputTokens
}
if contextInputTokens > 0 {
usage.InputTokens = contextInputTokens
}
if usage.TotalTokens == 0 {
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
}
@@ -935,7 +956,7 @@ func StreamEventStreamAsAnthropicWithContext(ctx context.Context, body io.Reader
"input_tokens": usage.InputTokens,
"output_tokens": usage.OutputTokens,
"cache_read_input_tokens": usage.CacheReadInputTokens,
"cache_creation_input_tokens": 0,
"cache_creation_input_tokens": usage.CacheCreationInputTokens,
},
}); err != nil {
return nil, err
@@ -1883,9 +1904,10 @@ func buildClaudeResponse(content string, toolUses []KiroToolUse, model string, u
"content": blocks,
"stop_reason": stopReason,
"usage": map[string]interface{}{
"input_tokens": usage.InputTokens,
"output_tokens": usage.OutputTokens,
"cache_read_input_tokens": usage.CacheReadInputTokens,
"input_tokens": usage.InputTokens,
"output_tokens": usage.OutputTokens,
"cache_read_input_tokens": usage.CacheReadInputTokens,
"cache_creation_input_tokens": usage.CacheCreationInputTokens,
},
}
result, _ := json.Marshal(response)
@@ -2649,6 +2671,9 @@ func updateUsageFromEvent(usage *Usage, eventType string, event map[string]inter
usage.InputTokens += value
}
}
if value, ok := firstInt(tokenUsage, "cacheCreationInputTokens", "cacheWriteInputTokens", "cacheCreateInputTokens", "cacheCreationTokens", "cacheWriteTokens", "uploadedInputTokens", "uploadInputTokens", "uploadedTokens", "uploadTokens"); ok {
usage.CacheCreationInputTokens = value
}
}
if value, ok := firstInt(event, "inputTokens", "inputTokenCount", "promptTokens", "prompt_tokens"); ok && value > 0 {
usage.InputTokens = value
@@ -2668,6 +2693,18 @@ func updateUsageFromEvent(usage *Usage, eventType string, event map[string]inter
if value, ok := firstInt(meta, "totalTokens", "totalTokenCount"); ok && value > 0 {
usage.TotalTokens = value
}
if value, ok := firstInt(event, "cacheReadInputTokens", "cachedInputTokens", "cacheReadTokens", "cachedTokens", "cached_tokens"); ok && value > 0 {
usage.CacheReadInputTokens = value
}
if value, ok := firstInt(meta, "cacheReadInputTokens", "cachedInputTokens", "cacheReadTokens", "cachedTokens", "cached_tokens"); ok && value > 0 {
usage.CacheReadInputTokens = value
}
if value, ok := firstInt(event, "cacheCreationInputTokens", "cacheWriteInputTokens", "cacheCreateInputTokens", "cacheCreationTokens", "cacheWriteTokens", "uploadedInputTokens", "uploadInputTokens", "uploadedTokens", "uploadTokens"); ok && value > 0 {
usage.CacheCreationInputTokens = value
}
if value, ok := firstInt(meta, "cacheCreationInputTokens", "cacheWriteInputTokens", "cacheCreateInputTokens", "cacheCreationTokens", "cacheWriteTokens", "uploadedInputTokens", "uploadInputTokens", "uploadedTokens", "uploadTokens"); ok && value > 0 {
usage.CacheCreationInputTokens = value
}
}
func firstInt(m map[string]interface{}, keys ...string) (int, bool) {
@@ -2679,6 +2716,39 @@ func firstInt(m map[string]interface{}, keys ...string) (int, bool) {
return 0, false
}
func firstFloat(m map[string]interface{}, keys ...string) (float64, bool) {
for _, key := range keys {
switch v := m[key].(type) {
case float64:
return v, true
case float32:
return float64(v), true
case int:
return float64(v), true
case int64:
return float64(v), true
case json.Number:
if parsed, err := v.Float64(); err == nil {
return parsed, true
}
case string:
if parsed, err := strconv.ParseFloat(strings.TrimSpace(v), 64); err == nil {
return parsed, true
}
}
}
return 0, false
}
func kiroContextWindowSize(model string) int {
switch MapModel(model) {
case "claude-opus-4.6", "claude-sonnet-4.6":
return 1_000_000
default:
return 200_000
}
}
func estimateKiroOutputTokens(content string, toolUses []KiroToolUse) int {
total := countKiroTextTokens(content)
for _, tool := range toolUses {
+105 -1
View File
@@ -292,6 +292,7 @@ func TestParseNonStreamingEventStreamUsageAliases(t *testing.T) {
"inputTokenCount": 12,
"completionTokens": 7,
"cachedTokens": 3,
"uploadedTokens": 5,
"totalTokenCount": 22,
},
},
@@ -302,8 +303,10 @@ func TestParseNonStreamingEventStreamUsageAliases(t *testing.T) {
require.Equal(t, 15, result.Usage.InputTokens)
require.Equal(t, 7, result.Usage.OutputTokens)
require.Equal(t, 3, result.Usage.CacheReadInputTokens)
require.Equal(t, 5, result.Usage.CacheCreationInputTokens)
require.Equal(t, 22, result.Usage.TotalTokens)
require.Equal(t, float64(3), gjson.GetBytes(result.ResponseBody, "usage.cache_read_input_tokens").Float())
require.Equal(t, float64(5), gjson.GetBytes(result.ResponseBody, "usage.cache_creation_input_tokens").Float())
}
func TestParseNonStreamingEventStreamEstimatesMissingOutputTokens(t *testing.T) {
@@ -316,7 +319,7 @@ func TestParseNonStreamingEventStreamEstimatesMissingOutputTokens(t *testing.T)
_, _ = stream.Write(buildEventStreamFrame(t, "messageMetadataEvent", map[string]any{
"messageMetadataEvent": map[string]any{
"tokenUsage": map[string]any{
"uncachedInputTokens": 12,
"uncachedInputTokens": 12,
"cacheReadInputTokens": 3,
},
},
@@ -581,6 +584,89 @@ func TestStreamEventStreamAsAnthropicSkipsLeadingWhitespaceOnlyChunk(t *testing.
require.NotContains(t, output, `"delta":{"text":"","type":"text_delta"}`)
}
func TestStreamEventStreamAsAnthropicUsageAliases(t *testing.T) {
stream := bytes.NewBuffer(nil)
_, _ = stream.Write(buildEventStreamFrame(t, "assistantResponseEvent", map[string]any{
"assistantResponseEvent": map[string]any{
"content": "hello",
},
}))
_, _ = stream.Write(buildEventStreamFrame(t, "metadataEvent", map[string]any{
"metadataEvent": map[string]any{
"tokenUsage": map[string]any{
"inputTokenCount": 12,
"completionTokens": 7,
"cachedTokens": 3,
"uploadedTokens": 5,
"totalTokenCount": 27,
},
},
}))
var out bytes.Buffer
result, err := StreamEventStreamAsAnthropic(context.Background(), stream, &out, "claude-sonnet-4-5", 9)
require.NoError(t, err)
require.Equal(t, 15, result.Usage.InputTokens)
require.Equal(t, 7, result.Usage.OutputTokens)
require.Equal(t, 3, result.Usage.CacheReadInputTokens)
require.Equal(t, 5, result.Usage.CacheCreationInputTokens)
require.Equal(t, 27, result.Usage.TotalTokens)
finalUsage := lastStreamUsage(t, out.String())
require.Equal(t, int64(15), finalUsage.Get("input_tokens").Int())
require.Equal(t, int64(7), finalUsage.Get("output_tokens").Int())
require.Equal(t, int64(3), finalUsage.Get("cache_read_input_tokens").Int())
require.Equal(t, int64(5), finalUsage.Get("cache_creation_input_tokens").Int())
}
func TestStreamEventStreamAsAnthropicEstimatesMissingOutputTokens(t *testing.T) {
stream := bytes.NewBuffer(nil)
_, _ = stream.Write(buildEventStreamFrame(t, "assistantResponseEvent", map[string]any{
"assistantResponseEvent": map[string]any{
"content": "Hello from Kiro",
},
}))
_, _ = stream.Write(buildEventStreamFrame(t, "messageMetadataEvent", map[string]any{
"messageMetadataEvent": map[string]any{
"tokenUsage": map[string]any{
"uncachedInputTokens": 9,
},
},
}))
var out bytes.Buffer
result, err := StreamEventStreamAsAnthropic(context.Background(), stream, &out, "claude-sonnet-4-5", 9)
require.NoError(t, err)
require.Equal(t, 9, result.Usage.InputTokens)
require.Equal(t, 4, result.Usage.OutputTokens)
require.Equal(t, 13, result.Usage.TotalTokens)
finalUsage := lastStreamUsage(t, out.String())
require.Equal(t, int64(4), finalUsage.Get("output_tokens").Int())
}
func TestStreamEventStreamAsAnthropicUsesContextUsageInputTokens(t *testing.T) {
stream := bytes.NewBuffer(nil)
_, _ = stream.Write(buildEventStreamFrame(t, "contextUsageEvent", map[string]any{
"contextUsageEvent": map[string]any{
"contextUsagePercentage": 2.5,
},
}))
_, _ = stream.Write(buildEventStreamFrame(t, "assistantResponseEvent", map[string]any{
"assistantResponseEvent": map[string]any{
"content": "ok",
},
}))
var out bytes.Buffer
result, err := StreamEventStreamAsAnthropic(context.Background(), stream, &out, "claude-opus-4.6", 9)
require.NoError(t, err)
require.Equal(t, 25000, result.Usage.InputTokens)
finalUsage := lastStreamUsage(t, out.String())
require.Equal(t, int64(25000), finalUsage.Get("input_tokens").Int())
}
func TestStreamEventStreamAsAnthropicSkipsTrailingWhitespaceOnlyChunk(t *testing.T) {
stream := bytes.NewBuffer(nil)
_, _ = stream.Write(buildEventStreamFrame(t, "assistantResponseEvent", map[string]any{
@@ -1271,3 +1357,21 @@ func buildEventStreamFrame(t *testing.T, eventType string, payload any) []byte {
require.NoError(t, binary.Write(frame, binary.BigEndian, uint32(0)))
return frame.Bytes()
}
func lastStreamUsage(t *testing.T, output string) gjson.Result {
t.Helper()
var usage gjson.Result
for _, frame := range strings.Split(output, "\n\n") {
if !strings.Contains(frame, "event: message_delta") {
continue
}
for _, line := range strings.Split(frame, "\n") {
if !strings.HasPrefix(line, "data: ") {
continue
}
usage = gjson.Get(strings.TrimPrefix(line, "data: "), "usage")
}
}
require.True(t, usage.Exists(), "message_delta usage not found in stream: %s", output)
return usage
}