feat: support explicit prompt token tracking in SSE parsing and stream handlers

internal/sse/parser.go

@@ -364,34 +364,50 @@ func hasContentFilterStatusValue(v any) bool {
 	return false
 }
 
-func extractAccumulatedTokenUsage(chunk map[string]any) int {
+func extractAccumulatedTokenUsage(chunk map[string]any) (int, int) {
 	return findAccumulatedTokenUsage(chunk)
 }
 
-func findAccumulatedTokenUsage(v any) int {
+func findAccumulatedTokenUsage(v any) (int, int) {
 	switch x := v.(type) {
 	case map[string]any:
 		if p, _ := x["p"].(string); strings.Contains(strings.ToLower(p), "accumulated_token_usage") {
 			if n, ok := toInt(x["v"]); ok && n > 0 {
-				return n
+				return 0, n
+			}
+		}
+		if p, _ := x["p"].(string); strings.Contains(strings.ToLower(p), "token_usage") {
+			if m, ok := x["v"].(map[string]any); ok {
+				p, _ := toInt(m["prompt_tokens"])
+				c, _ := toInt(m["completion_tokens"])
+				if p > 0 || c > 0 {
+					return p, c
+				}
 			}
 		}
 		if n, ok := toInt(x["accumulated_token_usage"]); ok && n > 0 {
-			return n
+			return 0, n
+		}
+		if usage, ok := x["token_usage"].(map[string]any); ok {
+			p, _ := toInt(usage["prompt_tokens"])
+			c, _ := toInt(usage["completion_tokens"])
+			if p > 0 || c > 0 {
+				return p, c
+			}
 		}
 		for _, vv := range x {
-			if n := findAccumulatedTokenUsage(vv); n > 0 {
-				return n
+			if p, c := findAccumulatedTokenUsage(vv); p > 0 || c > 0 {
+				return p, c
 			}
 		}
 	case []any:
 		for _, item := range x {
-			if n := findAccumulatedTokenUsage(item); n > 0 {
-				return n
+			if p, c := findAccumulatedTokenUsage(item); p > 0 || c > 0 {
+				return p, c
 			}
 		}
 	}
-	return 0
+	return 0, 0
 }
 
 func toInt(v any) (int, bool) {