fix(tokens): add stream options for usage inclusion

internal/llm/openai.go

@@ -35,12 +35,17 @@ type openAIMessage struct {
 	Content string `json:"content"`
 }
 
+type streamOptions struct {
+	IncludeUsage bool `json:"include_usage"`
+}
+
 type openAIReq struct {
-	Model       string          `json:"model"`
-	Messages    []openAIMessage `json:"messages"`
-	Stream      bool            `json:"stream"`
-	MaxTokens   int             `json:"max_tokens,omitempty"`
-	Temperature float64         `json:"temperature,omitempty"`
+	Model         string          `json:"model"`
+	Messages      []openAIMessage `json:"messages"`
+	Stream        bool            `json:"stream"`
+	MaxTokens     int             `json:"max_tokens,omitempty"`
+	Temperature   float64         `json:"temperature,omitempty"`
+	StreamOptions *streamOptions  `json:"stream_options,omitempty"`
 }
 
 type openAIResp struct {
@@ -199,11 +204,12 @@ func (c *OpenAIClient) GenerateResponse(ctx context.Context, req PromptRequest)
 
 func (c *OpenAIClient) StreamResponse(ctx context.Context, req PromptRequest, handler StreamHandler) (LLMResponse, error) {
 	body := openAIReq{
-		Model:       c.resolveModel(req.Model),
-		Messages:    c.buildMessages(req),
-		Stream:      true,
-		MaxTokens:   req.MaxTokens,
-		Temperature: req.Temperature,
+		Model:         c.resolveModel(req.Model),
+		Messages:      c.buildMessages(req),
+		Stream:        true,
+		MaxTokens:     req.MaxTokens,
+		Temperature:   req.Temperature,
+		StreamOptions: &streamOptions{IncludeUsage: true},
 	}
 	if body.MaxTokens <= 0 {
 		body.MaxTokens = 4096

internal/providers/openai.go

@@ -37,9 +37,10 @@ type openaiMessage struct {
 }
 
 type openaiRequest struct {
-	Model    string          `json:"model"`
-	Messages []openaiMessage `json:"messages"`
-	Stream   bool            `json:"stream"`
+	Model         string          `json:"model"`
+	Messages      []openaiMessage `json:"messages"`
+	Stream        bool            `json:"stream"`
+	StreamOptions *streamOptions  `json:"stream_options,omitempty"`
 }
 
 type openaiResponse struct {
@@ -63,6 +64,10 @@ type openaiDelta struct {
 	Content string `json:"content,omitempty"`
 }
 
+type streamOptions struct {
+	IncludeUsage bool `json:"include_usage"`
+}
+
 type openaiUsage struct {
 	PromptTokens     int `json:"prompt_tokens"`
 	CompletionTokens int `json:"completion_tokens"`
@@ -155,9 +160,10 @@ func (p *OpenAIProvider) ExecuteStream(ctx context.Context, req ai.Request) (io.
 	msgs := p.buildMessages(req)
 
 	body := openaiRequest{
-		Model:    model,
-		Messages: msgs,
-		Stream:   true,
+		Model:         model,
+		Messages:      msgs,
+		Stream:        true,
+		StreamOptions: &streamOptions{IncludeUsage: true},
 	}
 
 	payload, err := json.Marshal(body)
@@ -243,14 +249,14 @@ func (s *openaiSSEReader) Read(p []byte) (int, error) {
 			continue
 		}
 
-		if len(chunk.Choices) == 0 {
-			continue
-		}
-
 		if chunk.Usage != nil {
 			s.finalUsage = chunk.Usage
 		}
 
+		if len(chunk.Choices) == 0 {
+			continue
+		}
+
 		if chunk.Choices[0].Delta != nil && chunk.Choices[0].Delta.Content != "" {
 			n := copy(p, chunk.Choices[0].Delta.Content)
 			return n, nil

internal/providers/openrouter.go

@@ -110,9 +110,10 @@ func (p *OpenRouterProvider) ExecuteStream(ctx context.Context, req ai.Request)
 	msgs := p.buildMessages(req)
 
 	body := openrouterRequest{
-		Model:    model,
-		Messages: msgs,
-		Stream:   true,
+		Model:         model,
+		Messages:      msgs,
+		Stream:        true,
+		StreamOptions: &streamOptions{IncludeUsage: true},
 	}
 
 	payload, err := json.Marshal(body)
@@ -161,9 +162,10 @@ type openrouterMessage struct {
 }
 
 type openrouterRequest struct {
-	Model    string              `json:"model"`
-	Messages []openrouterMessage `json:"messages"`
-	Stream   bool                `json:"stream"`
+	Model         string              `json:"model"`
+	Messages      []openrouterMessage `json:"messages"`
+	Stream        bool                `json:"stream"`
+	StreamOptions *streamOptions      `json:"stream_options,omitempty"`
 }
 
 type openrouterResponse struct {
@@ -253,14 +255,14 @@ func (s *openrouterSSEReader) Read(p []byte) (int, error) {
 			continue
 		}
 
-		if len(chunk.Choices) == 0 {
-			continue
-		}
-
 		if chunk.Usage != nil {
 			s.finalUsage = chunk.Usage
 		}
 
+		if len(chunk.Choices) == 0 {
+			continue
+		}
+
 		if chunk.Choices[0].Delta != nil && chunk.Choices[0].Delta.Content != "" {
 			n := copy(p, chunk.Choices[0].Delta.Content)
 			return n, nil

internal/ui/stream.go

@@ -17,7 +17,6 @@ import (
 	"github.com/PizenLabs/izen/internal/modes"
 	"github.com/PizenLabs/izen/internal/modes/plan"
 	"github.com/PizenLabs/izen/internal/prompt"
-	"github.com/PizenLabs/izen/internal/providers"
 )
 
 // debugLogPayload writes the exact outgoing LLM payload to
@@ -185,8 +184,11 @@ func (m *model) streamCmd(content string) tea.Cmd {
 				streamCh <- tokenMsg(chunk)
 			}
 			if err == io.EOF {
-				if sr, ok := rawStream.(*providers.StreamResult); ok {
-					tokIn, tokOut = sr.Usage()
+				type usageProvider interface {
+					Usage() (input, output int)
+				}
+				if up, ok := rawStream.(usageProvider); ok {
+					tokIn, tokOut = up.Usage()
 				}
 				if tokIn == 0 && tokOut == 0 {
 					tokIn = len(content) / 4