1 files changed, 39 insertions, 18 deletions

diff --git a/internal/executor/gemini.go b/internal/executor/gemini.go
index d79c47d..a13321b 100644
--- a/internal/executor/gemini.go
+++ b/internal/executor/gemini.go
@@ -2,6 +2,7 @@ package executor
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"log/slog"
@@ -146,31 +147,51 @@ func parseGeminiStream(r io.Reader, w io.Writer, logger *slog.Logger) (float64,
 	}
 	logger.Debug("parseGeminiStream: raw output received", "output", string(fullOutput))
 
-	outputStr := strings.TrimSpace(string(fullOutput)) // Trim leading/trailing whitespace/newlines from the whole output
-
-	jsonContent := outputStr // Default to raw output if no markdown block is found or malformed
-	jsonStartIdx := strings.Index(outputStr, "```json")
-	if jsonStartIdx != -1 {
-		// Found "```json", now look for the closing "```"
-		jsonEndIdx := strings.LastIndex(outputStr, "```")
-		if jsonEndIdx != -1 && jsonEndIdx > jsonStartIdx {
-			// Extract content between the markdown fences.
-			jsonContent = outputStr[jsonStartIdx+len("```json"):jsonEndIdx]
-			jsonContent = strings.TrimSpace(jsonContent) // Trim again after extraction, to remove potential inner newlines
+	// Default: write raw content as-is (preserves trailing newline).
+	jsonContent := string(fullOutput)
+
+	// Unwrap markdown code fences if present.
+	trimmed := strings.TrimSpace(jsonContent)
+	if jsonStartIdx := strings.Index(trimmed, "```json"); jsonStartIdx != -1 {
+		if jsonEndIdx := strings.LastIndex(trimmed, "```"); jsonEndIdx != -1 && jsonEndIdx > jsonStartIdx {
+			inner := trimmed[jsonStartIdx+len("```json") : jsonEndIdx]
+			jsonContent = strings.TrimSpace(inner) + "\n"
 		} else {
-			logger.Warn("Malformed markdown JSON block from Gemini (missing closing ``` or invalid structure), falling back to raw output.", "outputLength", len(outputStr))
+			logger.Warn("malformed markdown JSON block from Gemini, falling back to raw output", "outputLength", len(jsonContent))
 		}
-	} else {
-		logger.Warn("No markdown JSON block found from Gemini, falling back to raw output.", "outputLength", len(outputStr))
 	}
 
-	// Write the (possibly extracted and trimmed) JSON content to the writer.
-	_, writeErr := w.Write([]byte(jsonContent))
-	if writeErr != nil {
+	// Write the (possibly extracted) JSON content to the writer.
+	if _, writeErr := w.Write([]byte(jsonContent)); writeErr != nil {
 		return 0, fmt.Errorf("writing extracted gemini json: %w", writeErr)
 	}
 
-	return 0, nil // For now, no cost/error parsing for Gemini stream
+	// Parse each line for result type to extract cost and execution errors.
+	var resultErr error
+	var costUSD float64
+	for _, line := range strings.Split(jsonContent, "\n") {
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+		var msg struct {
+			Type    string  `json:"type"`
+			IsError bool    `json:"is_error"`
+			Result  string  `json:"result"`
+			Cost    float64 `json:"total_cost_usd"`
+		}
+		if err := json.Unmarshal([]byte(line), &msg); err != nil {
+			continue
+		}
+		if msg.Type == "result" {
+			costUSD = msg.Cost
+			if msg.IsError {
+				resultErr = fmt.Errorf("gemini execution error: %s", msg.Result)
+			}
+		}
+	}
+
+	return costUSD, resultErr
 }
 
 func (r *GeminiRunner) buildArgs(t *task.Task, e *storage.Execution, questionFile string) []string {