diff --git a/server/pkg/agent/codebuddy.go b/server/pkg/agent/codebuddy.go index 328a1b63c..5ec76130a 100644 --- 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 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 +} diff --git a/server/pkg/agent/codebuddy_test.go b/server/pkg/agent/codebuddy_test.go index 170a656bb..57b448bac 100644 --- 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") + } +}