From f3355049bcacdc445e29625286e42fe34e30b341 Mon Sep 17 00:00:00 2001 From: Jiang Bohan Date: Tue, 14 Apr 2026 14:17:13 +0800 Subject: [PATCH] feat(agent): add live log support for Gemini CLI via stream-json MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch Gemini backend from `-o text` (batch output) to `-o stream-json` (NDJSON streaming) so tool calls, text, and errors are forwarded to the UI in real time instead of collected at the end. Parses all Gemini stream-json event types: init, message, tool_use, tool_result, error, and result — including per-model token usage from the result stats. --- server/pkg/agent/gemini.go | 200 +++++++++++++++++++++++++------- server/pkg/agent/gemini_test.go | 2 +- 2 files changed, 156 insertions(+), 46 deletions(-) diff --git a/server/pkg/agent/gemini.go b/server/pkg/agent/gemini.go index e34fffdba..1bde6f795 100644 --- a/server/pkg/agent/gemini.go +++ b/server/pkg/agent/gemini.go @@ -3,20 +3,15 @@ package agent import ( "bufio" "context" + "encoding/json" "fmt" - "io" "os/exec" "strings" "time" ) // geminiBackend implements Backend by spawning the Google Gemini CLI -// (`gemini -p --yolo -o text`) and collecting its stdout. -// -// This is a minimal v1 implementation — it captures the final text -// response but does not stream tool calls in real time. Follow-ups -// can move to `-o stream-json` and parse Gemini's event schema once -// we have a reliable reproduction of its output format. +// with `--output-format stream-json` and parsing its NDJSON event stream. type geminiBackend struct { cfg Config } @@ -59,10 +54,10 @@ func (b *geminiBackend) Execute(ctx context.Context, prompt string, opts ExecOpt b.cfg.Logger.Info("gemini started", "pid", cmd.Process.Pid, "cwd", opts.Cwd, "model", opts.Model) - msgCh := make(chan Message, 16) + msgCh := make(chan Message, 256) resCh := make(chan Result, 1) - // Close stdout when the context is cancelled so io.ReadAll unblocks. + // Close stdout when the context is cancelled so scanner.Scan() unblocks. go func() { <-runCtx.Done() _ = stdout.Close() @@ -74,66 +69,181 @@ func (b *geminiBackend) Execute(ctx context.Context, prompt string, opts ExecOpt defer close(resCh) startTime := time.Now() - output, readErr := io.ReadAll(bufio.NewReader(stdout)) + var output strings.Builder + var sessionID string + finalStatus := "completed" + var finalError string + usage := make(map[string]TokenUsage) - // Forward the full response as a single text message so the daemon - // can persist it verbatim. Tool streaming is intentionally omitted - // in v1; see the file-level comment. - text := strings.TrimRight(string(output), "\n") - if text != "" { - trySend(msgCh, Message{Type: MessageText, Content: text}) + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024) + + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + + var evt geminiStreamEvent + if err := json.Unmarshal([]byte(line), &evt); err != nil { + continue + } + + switch evt.Type { + case "init": + sessionID = evt.SessionID + trySend(msgCh, Message{Type: MessageStatus, Status: "running"}) + + case "message": + if evt.Role == "assistant" && evt.Content != "" { + output.WriteString(evt.Content) + trySend(msgCh, Message{Type: MessageText, Content: evt.Content}) + } + + case "tool_use": + var params map[string]any + if evt.Parameters != nil { + _ = json.Unmarshal(evt.Parameters, ¶ms) + } + trySend(msgCh, Message{ + Type: MessageToolUse, + Tool: evt.ToolName, + CallID: evt.ToolID, + Input: params, + }) + + case "tool_result": + trySend(msgCh, Message{ + Type: MessageToolResult, + CallID: evt.ToolID, + Output: evt.Output, + }) + + case "error": + trySend(msgCh, Message{ + Type: MessageError, + Content: evt.Message, + }) + + case "result": + if evt.Status == "error" && evt.Error != nil { + finalStatus = "failed" + finalError = evt.Error.Message + } + if evt.Stats != nil { + b.accumulateUsage(usage, evt.Stats) + } + } } waitErr := cmd.Wait() - durationMs := time.Since(startTime).Milliseconds() + duration := time.Since(startTime) - result := Result{ - Status: "completed", - Output: text, - DurationMs: durationMs, - } - - // Check context errors first — timeout and cancellation kill the - // process, which can cause readErr/waitErr as side effects. The - // context error is the root cause and determines the correct status. if runCtx.Err() == context.DeadlineExceeded { - result.Status = "timeout" - result.Error = fmt.Sprintf("gemini timed out after %s", timeout) + finalStatus = "timeout" + finalError = fmt.Sprintf("gemini timed out after %s", timeout) } else if runCtx.Err() == context.Canceled { - result.Status = "aborted" - result.Error = "execution cancelled" - } else if readErr != nil { - result.Status = "failed" - result.Error = fmt.Sprintf("read stdout: %s", readErr.Error()) - } else if waitErr != nil { - result.Status = "failed" - result.Error = waitErr.Error() + finalStatus = "aborted" + finalError = "execution cancelled" + } else if waitErr != nil && finalStatus == "completed" { + finalStatus = "failed" + finalError = fmt.Sprintf("gemini exited with error: %v", waitErr) } - resCh <- result + b.cfg.Logger.Info("gemini finished", "pid", cmd.Process.Pid, "status", finalStatus, "duration", duration.Round(time.Millisecond).String()) + + resCh <- Result{ + Status: finalStatus, + Output: output.String(), + Error: finalError, + DurationMs: duration.Milliseconds(), + SessionID: sessionID, + Usage: usage, + } }() return &Session{Messages: msgCh, Result: resCh}, nil } +// accumulateUsage extracts per-model token usage from Gemini's result stats. +func (b *geminiBackend) accumulateUsage(usage map[string]TokenUsage, stats *geminiStreamStats) { + for model, m := range stats.Models { + u := usage[model] + u.InputTokens += int64(m.InputTokens) + u.OutputTokens += int64(m.OutputTokens) + u.CacheReadTokens += int64(m.Cached) + usage[model] = u + } +} + +// ── Gemini stream-json event types ── + +type geminiStreamEvent struct { + Type string `json:"type"` + Timestamp string `json:"timestamp,omitempty"` + SessionID string `json:"session_id,omitempty"` + Model string `json:"model,omitempty"` + + // message fields + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + Delta bool `json:"delta,omitempty"` + + // tool_use fields + ToolName string `json:"tool_name,omitempty"` + ToolID string `json:"tool_id,omitempty"` + Parameters json.RawMessage `json:"parameters,omitempty"` + + // tool_result fields + Status string `json:"status,omitempty"` + Output string `json:"output,omitempty"` + + // error fields + Severity string `json:"severity,omitempty"` + Message string `json:"message,omitempty"` + + // result fields + Error *geminiStreamError `json:"error,omitempty"` + Stats *geminiStreamStats `json:"stats,omitempty"` +} + +type geminiStreamError struct { + Type string `json:"type"` + Message string `json:"message"` +} + +type geminiStreamStats struct { + TotalTokens int `json:"total_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + DurationMs int `json:"duration_ms"` + ToolCalls int `json:"tool_calls"` + Models map[string]geminiModelStats `json:"models,omitempty"` +} + +type geminiModelStats struct { + TotalTokens int `json:"total_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + Cached int `json:"cached"` +} + +// ── Arg builder ── + // buildGeminiArgs assembles the argv for a one-shot gemini invocation. // // Flags: // // -p / --prompt non-interactive prompt (the user's task) -// --yolo auto-approve all tool executions (equivalent to -// claude's --permission-mode bypassPermissions) -// -o text plain text output (stream-json is a follow-up) -// -m optional model override (from MULTICA_GEMINI_MODEL) +// --yolo auto-approve all tool executions +// -o stream-json streaming NDJSON output for live events +// -m optional model override // -r resume a previous session (if provided) -// -// Note: gemini reads stdin and appends it to -p when both are present. -// The daemon does not pipe stdin, so the prompt comes exclusively from -p. func buildGeminiArgs(prompt string, opts ExecOptions) []string { args := []string{ "-p", prompt, "--yolo", - "-o", "text", + "-o", "stream-json", } if opts.Model != "" { args = append(args, "-m", opts.Model) diff --git a/server/pkg/agent/gemini_test.go b/server/pkg/agent/gemini_test.go index 8fc4e6339..3a95bdd75 100644 --- a/server/pkg/agent/gemini_test.go +++ b/server/pkg/agent/gemini_test.go @@ -11,7 +11,7 @@ func TestBuildGeminiArgsBaseline(t *testing.T) { expected := []string{ "-p", "write a haiku", "--yolo", - "-o", "text", + "-o", "stream-json", } if len(args) != len(expected) {