feat(agent): add live log support for Gemini CLI via stream-json

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.
This commit is contained in:
Jiang Bohan
2026-04-14 14:17:13 +08:00
parent dca86acc69
commit f3355049bc
2 changed files with 156 additions and 46 deletions

View File

@@ -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 <prompt> --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, &params)
}
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 <model> optional model override (from MULTICA_GEMINI_MODEL)
// --yolo auto-approve all tool executions
// -o stream-json streaming NDJSON output for live events
// -m <model> optional model override
// -r <session> 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)

View File

@@ -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) {