feat(agent): implement codebuddyBackend.Execute with stream-json parsing

Co-authored-by: multica-agent <github@multica.ai>
2026-07-05 13:29:44 +02:00 · 2026-05-25 11:31:50 +08:00
parent 160a62ca9f
commit d2aef41fbf
2 changed files with 637 additions and 0 deletions
--- a/server/pkg/agent/codebuddy.go
+++ b/server/pkg/agent/codebuddy.go
@@ -1,8 +1,17 @@
 package agent

 import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
+	"io"
 	"log/slog"
+	"os"
+	"os/exec"
+	"strings"
+	"time"
 )

 // codebuddyBackend implements Backend by spawning the Claude Code CLI
@@ -54,3 +63,392 @@ func buildCodebuddyArgs(opts ExecOptions, logger *slog.Logger) []string {
 	args = append(args, filterCustomArgs(opts.CustomArgs, codebuddyBlockedArgs, logger)...)
 	return args
 }
+
+func (b *codebuddyBackend) Execute(ctx context.Context, prompt string, opts ExecOptions) (*Session, error) {
+	execPath := b.cfg.ExecutablePath
+	if execPath == "" {
+		execPath = "codebuddy"
+	}
+	if _, err := exec.LookPath(execPath); err != nil {
+		return nil, fmt.Errorf("codebuddy executable not found at %q: %w", execPath, err)
+	}
+
+	timeout := opts.Timeout
+	if timeout == 0 {
+		timeout = 20 * time.Minute
+	}
+	runCtx, cancel := context.WithTimeout(ctx, timeout)
+
+	args := buildCodebuddyArgs(opts, b.cfg.Logger)
+
+	// If the caller provided an MCP config, write it to a temp file and pass
+	// --mcp-config <path> so the agent uses a controlled set of MCP servers.
+	var mcpConfigPath string
+	var mcpFileCleanup func()
+	if len(opts.McpConfig) > 0 {
+		path, err := writeMcpConfigToTemp(opts.McpConfig)
+		if err != nil {
+			cancel()
+			return nil, err
+		}
+		mcpConfigPath = path
+		mcpFileCleanup = func() { os.Remove(mcpConfigPath) }
+		args = append(args, "--mcp-config", mcpConfigPath)
+	}
+	// Clean up the temp file if we return before the goroutine takes ownership.
+	defer func() {
+		if mcpFileCleanup != nil {
+			mcpFileCleanup()
+		}
+	}()
+
+	cmd := exec.CommandContext(runCtx, execPath, args...)
+	hideAgentWindow(cmd)
+	b.cfg.Logger.Info("agent command", "exec", execPath, "args", args)
+	cmd.WaitDelay = 10 * time.Second
+	if opts.Cwd != "" {
+		cmd.Dir = opts.Cwd
+	}
+	cmd.Env = buildEnv(b.cfg.Env)
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("codebuddy stdout pipe: %w", err)
+	}
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("codebuddy stdin pipe: %w", err)
+	}
+	closeStdin := func() {
+		if stdin != nil {
+			_ = stdin.Close()
+			stdin = nil
+		}
+	}
+
+	stderrBuf := newStderrTail(newLogWriter(b.cfg.Logger, "[codebuddy:stderr] "), agentStderrTailBytes)
+	cmd.Stderr = stderrBuf
+
+	if err := cmd.Start(); err != nil {
+		closeStdin()
+		cancel()
+		return nil, fmt.Errorf("start codebuddy: %w", err)
+	}
+	if err := writeCodebuddyInput(stdin, prompt); err != nil {
+		closeStdin()
+		cancel()
+		_ = cmd.Wait()
+		return nil, errors.New(withAgentStderr(fmt.Sprintf("write codebuddy input: %v", err), "codebuddy", stderrBuf.Tail()))
+	}
+	closeStdin()
+
+	b.cfg.Logger.Info("codebuddy started", "pid", cmd.Process.Pid, "cwd", opts.Cwd, "model", opts.Model)
+
+	// cmd.Start() succeeded — transfer temp file ownership to the goroutine.
+	mcpFileCleanup = nil
+
+	msgCh := make(chan Message, 256)
+	resCh := make(chan Result, 1)
+
+	go func() {
+		defer cancel()
+		defer close(msgCh)
+		defer close(resCh)
+		if mcpConfigPath != "" {
+			defer os.Remove(mcpConfigPath)
+		}
+
+		startTime := time.Now()
+		var output strings.Builder
+		var sessionID string
+		finalStatus := "completed"
+		var finalError string
+		usage := make(map[string]TokenUsage)
+
+		// Close stdout when the context is cancelled so scanner.Scan() unblocks.
+		go func() {
+			<-runCtx.Done()
+			_ = stdout.Close()
+		}()
+
+		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 msg codebuddySDKMessage
+			if err := json.Unmarshal([]byte(line), &msg); err != nil {
+				continue
+			}
+
+			switch msg.Type {
+			case "assistant":
+				b.handleAssistant(msg, msgCh, &output, usage)
+			case "user":
+				b.handleUser(msg, msgCh)
+			case "system":
+				if msg.SessionID != "" {
+					sessionID = msg.SessionID
+				}
+				trySend(msgCh, Message{Type: MessageStatus, Status: "running", SessionID: sessionID})
+			case "result":
+				closeStdin()
+				sessionID = msg.SessionID
+				if msg.ResultText != "" {
+					output.Reset()
+					output.WriteString(msg.ResultText)
+				}
+				if resultUsage := codebuddyResultUsage(msg, opts.Model); len(resultUsage) > 0 {
+					usage = resultUsage
+				}
+				if msg.IsError {
+					finalStatus = "failed"
+					finalError = msg.ResultText
+				}
+			case "log":
+				if msg.Log != nil {
+					trySend(msgCh, Message{
+						Type:    MessageLog,
+						Level:   msg.Log.Level,
+						Content: msg.Log.Message,
+					})
+				}
+			}
+		}
+
+		// Wait for process exit.
+		exitErr := cmd.Wait()
+		duration := time.Since(startTime)
+
+		if runCtx.Err() == context.DeadlineExceeded {
+			finalStatus = "timeout"
+			finalError = fmt.Sprintf("codebuddy timed out after %s", timeout)
+		} else if runCtx.Err() == context.Canceled {
+			finalStatus = "aborted"
+			finalError = "execution cancelled"
+		} else if exitErr != nil && finalStatus == "completed" {
+			finalStatus = "failed"
+			finalError = fmt.Sprintf("codebuddy exited with error: %v", exitErr)
+		}
+
+		if finalError != "" {
+			finalError = withAgentStderr(finalError, "codebuddy", stderrBuf.Tail())
+		}
+
+		b.cfg.Logger.Info("codebuddy finished", "pid", cmd.Process.Pid, "status", finalStatus, "duration", duration.Round(time.Millisecond).String())
+
+		reportedSessionID := resolveSessionID(opts.ResumeSessionID, sessionID, finalStatus == "failed")
+		if reportedSessionID != sessionID {
+			b.cfg.Logger.Info("codebuddy resume did not land; clearing fresh session id for daemon fallback",
+				"requested_resume", opts.ResumeSessionID,
+				"emitted_session", sessionID,
+			)
+		}
+
+		resCh <- Result{
+			Status:     finalStatus,
+			Output:     output.String(),
+			Error:      finalError,
+			DurationMs: duration.Milliseconds(),
+			SessionID:  reportedSessionID,
+			Usage:      usage,
+		}
+	}()
+
+	return &Session{Messages: msgCh, Result: resCh}, nil
+}
+
+func (b *codebuddyBackend) handleAssistant(msg codebuddySDKMessage, ch chan<- Message, output *strings.Builder, usage map[string]TokenUsage) {
+	var content codebuddyMessageContent
+	if err := json.Unmarshal(msg.Message, &content); err != nil {
+		return
+	}
+
+	// Accumulate token usage per model.
+	if content.Usage != nil && content.Model != "" {
+		u := usage[content.Model]
+		u.InputTokens += content.Usage.InputTokens
+		u.OutputTokens += content.Usage.OutputTokens
+		u.CacheReadTokens += content.Usage.CacheReadInputTokens
+		u.CacheWriteTokens += content.Usage.CacheCreationInputTokens
+		usage[content.Model] = u
+	}
+
+	for _, block := range content.Content {
+		switch block.Type {
+		case "text":
+			if block.Text != "" {
+				output.WriteString(block.Text)
+				trySend(ch, Message{Type: MessageText, Content: block.Text})
+			}
+		case "thinking":
+			if block.Text != "" {
+				trySend(ch, Message{Type: MessageThinking, Content: block.Text})
+			}
+		case "tool_use":
+			var input map[string]any
+			if block.Input != nil {
+				_ = json.Unmarshal(block.Input, &input)
+			}
+			trySend(ch, Message{
+				Type:   MessageToolUse,
+				Tool:   block.Name,
+				CallID: block.ID,
+				Input:  input,
+			})
+		}
+	}
+}
+
+func (b *codebuddyBackend) handleUser(msg codebuddySDKMessage, ch chan<- Message) {
+	var content codebuddyMessageContent
+	if err := json.Unmarshal(msg.Message, &content); err != nil {
+		return
+	}
+
+	for _, block := range content.Content {
+		if block.Type == "tool_result" {
+			resultStr := ""
+			if block.Content != nil {
+				resultStr = string(block.Content)
+			}
+			trySend(ch, Message{
+				Type:   MessageToolResult,
+				CallID: block.ToolUseID,
+				Output: resultStr,
+			})
+		}
+	}
+}
+
+func writeCodebuddyInput(w io.Writer, prompt string) error {
+	payload := map[string]any{
+		"type": "user",
+		"message": map[string]any{
+			"role": "user",
+			"content": []map[string]string{
+				{
+					"type": "text",
+					"text": prompt,
+				},
+			},
+		},
+	}
+	data, err := json.Marshal(payload)
+	if err != nil {
+		return fmt.Errorf("marshal codebuddy input: %w", err)
+	}
+	data = append(data, '\n')
+	if _, err := w.Write(data); err != nil {
+		return err
+	}
+	return nil
+}
+
+// ── Codebuddy SDK JSON types ──
+
+type codebuddySDKMessage struct {
+	Type      string          `json:"type"`
+	Message   json.RawMessage `json:"message,omitempty"`
+	Subtype   string          `json:"subtype,omitempty"`
+	SessionID string          `json:"session_id,omitempty"`
+	Model     string          `json:"model,omitempty"`
+
+	// result fields
+	ResultText string                               `json:"result,omitempty"`
+	IsError    bool                                  `json:"is_error,omitempty"`
+	DurationMs float64                              `json:"duration_ms,omitempty"`
+	NumTurns   int                                  `json:"num_turns,omitempty"`
+	Usage      *codebuddyUsage                      `json:"usage,omitempty"`
+	ModelUsage map[string]codebuddyResultModelUsage `json:"modelUsage,omitempty"`
+
+	// log fields
+	Log *codebuddyLogEntry `json:"log,omitempty"`
+}
+
+type codebuddyLogEntry struct {
+	Level   string `json:"level"`
+	Message string `json:"message"`
+}
+
+type codebuddyMessageContent struct {
+	Role    string                 `json:"role"`
+	Model   string                 `json:"model"`
+	Content []codebuddyContentBlock `json:"content"`
+	Usage   *codebuddyUsage        `json:"usage,omitempty"`
+}
+
+type codebuddyUsage struct {
+	InputTokens              int64 `json:"input_tokens"`
+	OutputTokens             int64 `json:"output_tokens"`
+	CacheReadInputTokens     int64 `json:"cache_read_input_tokens"`
+	CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
+}
+
+type codebuddyResultModelUsage struct {
+	InputTokens              int64 `json:"inputTokens"`
+	OutputTokens             int64 `json:"outputTokens"`
+	CacheReadInputTokens     int64 `json:"cacheReadInputTokens"`
+	CacheCreationInputTokens int64 `json:"cacheCreationInputTokens"`
+}
+
+type codebuddyContentBlock struct {
+	Type      string          `json:"type"`
+	Text      string          `json:"text,omitempty"`
+	ID        string          `json:"id,omitempty"`
+	Name      string          `json:"name,omitempty"`
+	Input     json.RawMessage `json:"input,omitempty"`
+	ToolUseID string          `json:"tool_use_id,omitempty"`
+	Content   json.RawMessage `json:"content,omitempty"`
+}
+
+func codebuddyResultUsage(msg codebuddySDKMessage, fallbackModel string) map[string]TokenUsage {
+	if len(msg.ModelUsage) > 0 {
+		usage := make(map[string]TokenUsage, len(msg.ModelUsage))
+		for model, u := range msg.ModelUsage {
+			if model == "" || !codebuddyUsageHasTokens(u.InputTokens, u.OutputTokens, u.CacheReadInputTokens, u.CacheCreationInputTokens) {
+				continue
+			}
+			usage[model] = TokenUsage{
+				InputTokens:      u.InputTokens,
+				OutputTokens:     u.OutputTokens,
+				CacheReadTokens:  u.CacheReadInputTokens,
+				CacheWriteTokens: u.CacheCreationInputTokens,
+			}
+		}
+		if len(usage) > 0 {
+			return usage
+		}
+	}
+
+	model := msg.Model
+	if model == "" {
+		model = fallbackModel
+	}
+	if msg.Usage == nil || model == "" || !codebuddyUsageHasTokens(
+		msg.Usage.InputTokens,
+		msg.Usage.OutputTokens,
+		msg.Usage.CacheReadInputTokens,
+		msg.Usage.CacheCreationInputTokens,
+	) {
+		return nil
+	}
+	return map[string]TokenUsage{
+		model: {
+			InputTokens:      msg.Usage.InputTokens,
+			OutputTokens:     msg.Usage.OutputTokens,
+			CacheReadTokens:  msg.Usage.CacheReadInputTokens,
+			CacheWriteTokens: msg.Usage.CacheCreationInputTokens,
+		},
+	}
+}
+
+func codebuddyUsageHasTokens(input, output, cacheRead, cacheWrite int64) bool {
+	return input > 0 || output > 0 || cacheRead > 0 || cacheWrite > 0
+}
--- a/server/pkg/agent/codebuddy_test.go
+++ b/server/pkg/agent/codebuddy_test.go
@@ -1,9 +1,14 @@
 package agent

 import (
+	"context"
+	"encoding/json"
 	"log/slog"
+	"path/filepath"
+	"runtime"
 	"strings"
 	"testing"
+	"time"
 )

 func TestBuildCodebuddyArgs_Basic(t *testing.T) {
@@ -138,3 +143,237 @@ func TestBuildCodebuddyArgs_Resume(t *testing.T) {
 		t.Fatalf("expected --resume sess-abc123 in args: %v", args)
 	}
 }
+
+func TestCodebuddyExecute_Success(t *testing.T) {
+	t.Parallel()
+	if runtime.GOOS == "windows" {
+		t.Skip("shell-script fixture is POSIX-only")
+	}
+
+	fakePath := filepath.Join(t.TempDir(), "codebuddy")
+	script := "#!/bin/sh\n" +
+		"cat >/dev/null\n" +
+		`printf '%s\n' '{"type":"system","session_id":"sess-cb-001"}'` + "\n" +
+		`printf '%s\n' '{"type":"assistant","message":{"role":"assistant","model":"claude-sonnet-4-20250514","content":[{"type":"text","text":"Hello from codebuddy"}]}}'` + "\n" +
+		`printf '%s\n' '{"type":"result","subtype":"success","is_error":false,"session_id":"sess-cb-001","result":"Hello from codebuddy","modelUsage":{"claude-sonnet-4-20250514":{"inputTokens":100,"outputTokens":50,"cacheReadInputTokens":10,"cacheCreationInputTokens":5}}}'` + "\n"
+	writeTestExecutable(t, fakePath, []byte(script))
+
+	b := &codebuddyBackend{cfg: Config{ExecutablePath: fakePath, Logger: slog.Default()}}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	session, err := b.Execute(ctx, "say hello", ExecOptions{Timeout: 5 * time.Second})
+	if err != nil {
+		t.Fatalf("execute: %v", err)
+	}
+
+	// Drain messages.
+	var gotText bool
+	for msg := range session.Messages {
+		if msg.Type == MessageText && msg.Content == "Hello from codebuddy" {
+			gotText = true
+		}
+	}
+	if !gotText {
+		t.Fatal("expected text message 'Hello from codebuddy'")
+	}
+
+	select {
+	case result, ok := <-session.Result:
+		if !ok {
+			t.Fatal("result channel closed without a value")
+		}
+		if result.Status != "completed" {
+			t.Fatalf("expected status=completed, got %q (error=%q)", result.Status, result.Error)
+		}
+		if result.Output != "Hello from codebuddy" {
+			t.Fatalf("expected output 'Hello from codebuddy', got %q", result.Output)
+		}
+		if result.SessionID != "sess-cb-001" {
+			t.Fatalf("expected session_id=sess-cb-001, got %q", result.SessionID)
+		}
+		usage, ok := result.Usage["claude-sonnet-4-20250514"]
+		if !ok {
+			t.Fatalf("expected usage for claude-sonnet-4-20250514, got %#v", result.Usage)
+		}
+		if usage.InputTokens != 100 || usage.OutputTokens != 50 || usage.CacheReadTokens != 10 || usage.CacheWriteTokens != 5 {
+			t.Fatalf("unexpected usage: %+v", usage)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("timeout waiting for result")
+	}
+}
+
+func TestCodebuddyExecute_NotFound(t *testing.T) {
+	t.Parallel()
+
+	b := &codebuddyBackend{cfg: Config{ExecutablePath: "/nonexistent/path/codebuddy", Logger: slog.Default()}}
+
+	ctx := context.Background()
+	_, err := b.Execute(ctx, "prompt", ExecOptions{})
+	if err == nil {
+		t.Fatal("expected error for missing executable")
+	}
+	if !strings.Contains(err.Error(), "codebuddy executable not found") {
+		t.Fatalf("expected 'codebuddy executable not found' in error, got %q", err.Error())
+	}
+}
+
+func TestCodebuddyExecuteSurfacesStderr(t *testing.T) {
+	t.Parallel()
+	if runtime.GOOS == "windows" {
+		t.Skip("shell-script fixture is POSIX-only")
+	}
+
+	fakePath := filepath.Join(t.TempDir(), "codebuddy")
+	script := "#!/bin/sh\n" +
+		"cat >/dev/null\n" +
+		"echo \"FATAL ERROR: segfault in codebuddy runtime\" >&2\n" +
+		"exit 1\n"
+	writeTestExecutable(t, fakePath, []byte(script))
+
+	b := &codebuddyBackend{cfg: Config{ExecutablePath: fakePath, Logger: slog.Default()}}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	session, err := b.Execute(ctx, "prompt-ignored", ExecOptions{Timeout: 5 * time.Second})
+	if err != nil {
+		t.Fatalf("execute: %v", err)
+	}
+
+	// Drain messages.
+	go func() {
+		for range session.Messages {
+		}
+	}()
+
+	select {
+	case result, ok := <-session.Result:
+		if !ok {
+			t.Fatal("result channel closed without a value")
+		}
+		if result.Status != "failed" {
+			t.Fatalf("expected status=failed, got %q (error=%q)", result.Status, result.Error)
+		}
+		if !strings.Contains(result.Error, "codebuddy exited with error") {
+			t.Fatalf("expected error to mention exit, got %q", result.Error)
+		}
+		if !strings.Contains(result.Error, "segfault in codebuddy runtime") {
+			t.Fatalf("expected error to include stderr content, got %q", result.Error)
+		}
+		if !strings.Contains(result.Error, "codebuddy stderr:") {
+			t.Fatalf("expected stderr label in error, got %q", result.Error)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("timeout waiting for result")
+	}
+}
+
+func TestWriteCodebuddyInput(t *testing.T) {
+	t.Parallel()
+
+	var buf strings.Builder
+	err := writeCodebuddyInput(&buf, "hello world")
+	if err != nil {
+		t.Fatalf("writeCodebuddyInput: %v", err)
+	}
+
+	data := buf.String()
+	if len(data) == 0 || data[len(data)-1] != '\n' {
+		t.Fatalf("expected newline-terminated payload, got %q", data)
+	}
+
+	var payload map[string]any
+	if err := json.Unmarshal([]byte(strings.TrimSpace(data)), &payload); err != nil {
+		t.Fatalf("unmarshal payload: %v", err)
+	}
+	if payload["type"] != "user" {
+		t.Fatalf("expected type user, got %v", payload["type"])
+	}
+
+	message, ok := payload["message"].(map[string]any)
+	if !ok {
+		t.Fatalf("expected message object, got %T", payload["message"])
+	}
+	if message["role"] != "user" {
+		t.Fatalf("expected role user, got %v", message["role"])
+	}
+
+	content, ok := message["content"].([]any)
+	if !ok || len(content) != 1 {
+		t.Fatalf("expected one content block, got %v", message["content"])
+	}
+	block, ok := content[0].(map[string]any)
+	if !ok {
+		t.Fatalf("expected content block object, got %T", content[0])
+	}
+	if block["type"] != "text" || block["text"] != "hello world" {
+		t.Fatalf("unexpected content block: %v", block)
+	}
+}
+
+func TestCodebuddyHandleAssistantText(t *testing.T) {
+	t.Parallel()
+
+	b := &codebuddyBackend{cfg: Config{Logger: slog.Default()}}
+	ch := make(chan Message, 10)
+	var output strings.Builder
+
+	msg := codebuddySDKMessage{
+		Type: "assistant",
+		Message: mustMarshal(t, codebuddyMessageContent{
+			Role: "assistant",
+			Content: []codebuddyContentBlock{
+				{Type: "text", Text: "codebuddy says hi"},
+			},
+		}),
+	}
+
+	b.handleAssistant(msg, ch, &output, make(map[string]TokenUsage))
+
+	if output.String() != "codebuddy says hi" {
+		t.Fatalf("expected output 'codebuddy says hi', got %q", output.String())
+	}
+	select {
+	case m := <-ch:
+		if m.Type != MessageText || m.Content != "codebuddy says hi" {
+			t.Fatalf("unexpected message: %+v", m)
+		}
+	default:
+		t.Fatal("expected message on channel")
+	}
+}
+
+func TestCodebuddyHandleUserToolResult(t *testing.T) {
+	t.Parallel()
+
+	b := &codebuddyBackend{cfg: Config{Logger: slog.Default()}}
+	ch := make(chan Message, 10)
+
+	msg := codebuddySDKMessage{
+		Type: "user",
+		Message: mustMarshal(t, codebuddyMessageContent{
+			Role: "user",
+			Content: []codebuddyContentBlock{
+				{
+					Type:      "tool_result",
+					ToolUseID: "call-cb-1",
+					Content:   mustMarshal(t, "tool output here"),
+				},
+			},
+		}),
+	}
+
+	b.handleUser(msg, ch)
+
+	select {
+	case m := <-ch:
+		if m.Type != MessageToolResult || m.CallID != "call-cb-1" {
+			t.Fatalf("unexpected message: %+v", m)
+		}
+	default:
+		t.Fatal("expected message on channel")
+	}
+}