agents_and_robots/shell/llm/claudecode.go

package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"github.com/enmanuel/agents/internal/config"
	coretypes "github.com/enmanuel/agents/pkg/llm"
)

const (
	defaultClaudeBinary  = "claude"
	defaultClaudeTimeout = 5 * time.Minute
)

// claudeJSONOutput represents the JSON output from `claude -p --output-format json`.
type claudeJSONOutput struct {
	Type         string           `json:"type"`
	Subtype      string           `json:"subtype"`
	CostUSD      float64          `json:"cost_usd"`
	IsError      bool             `json:"is_error"`
	Duration     float64          `json:"duration_api_ms"`
	NumTurns     int              `json:"num_turns"`
	Result       string           `json:"result"`
	SessionID    string           `json:"session_id"`
	TotalCost    float64          `json:"total_cost_usd"`
	Usage        claudeUsage      `json:"usage"`
	ContentBlock []claudeContent  `json:"content"`
}

type claudeUsage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

type claudeContent struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

// NewClaudeCodeComplete creates a CompleteFunc that executes `claude -p` as a subprocess.
func NewClaudeCodeComplete(cfg config.ClaudeCodeCfg, log *slog.Logger) coretypes.CompleteFunc {
	binary := cfg.Binary
	if binary == "" {
		binary = defaultClaudeBinary
	}

	timeout := cfg.Timeout
	if timeout <= 0 {
		timeout = defaultClaudeTimeout
	}

	// Resolve working directory once at init time.
	workDir := cfg.WorkingDir
	if workDir == "" {
		tmp, err := os.MkdirTemp("", "claude-agent-*")
		if err != nil {
			log.Error("claude-code: failed to create temp working dir", "err", err)
			// Fall through — cmd.Dir will remain empty (inherits CWD).
		} else {
			workDir = tmp
			log.Warn("claude-code working_dir is empty, using temporary directory",
				"dir", workDir,
			)
		}
	} else {
		// Ensure configured directory exists.
		if err := os.MkdirAll(workDir, 0o755); err != nil {
			log.Error("claude-code: failed to create working dir", "dir", workDir, "err", err)
		}
	}

	return func(ctx context.Context, req coretypes.CompletionRequest) (coretypes.CompletionResponse, error) {
		ctx, cancel := context.WithTimeout(ctx, timeout)
		defer cancel()

		args := buildClaudeArgs(cfg, req)

		prompt := flattenMessages(req.Messages)

		log.Debug("claude_code_exec",
			"binary", binary,
			"args", strings.Join(args, " "),
			"prompt_len", len(prompt),
			"working_dir", workDir,
		)

		cmd := exec.CommandContext(ctx, binary, args...)
		if workDir != "" {
			cmd.Dir = workDir
		}
		// Build clean env: inherit parent but remove ANTHROPIC_API_KEY
		// so claude uses its own OAuth auth instead of a potentially invalid key.
		cmd.Env = filterEnv(os.Environ(), "ANTHROPIC_API_KEY")
		cmd.Stdin = strings.NewReader(prompt)

		// Create a new process group so we can kill claude + all its children.
		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
		// Override the default cancel behavior: kill the entire process group
		// instead of just the main process, preventing orphaned child processes.
		cmd.Cancel = func() error {
			if cmd.Process != nil {
				pgid := cmd.Process.Pid
				log.Info("killing claude-code process group", "pgid", pgid)
				// Negative PID = kill entire process group
				return syscall.Kill(-pgid, syscall.SIGKILL)
			}
			return nil
		}

		var stdout, stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr

		start := time.Now()
		err := cmd.Run()
		elapsed := time.Since(start)

		// Ensure the process group is fully dead after Run returns,
		// even if cmd.Run() returned without triggering Cancel (normal exit).
		if cmd.Process != nil {
			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}

		log.Debug("claude_code_done",
			"elapsed_ms", elapsed.Milliseconds(),
			"stdout_len", stdout.Len(),
			"stderr_len", stderr.Len(),
			"exit_err", err,
		)

		return parseClaudeOutput(stdout.Bytes(), stderr.Bytes(), err, elapsed, log)
	}
}

// buildClaudeArgs constructs the CLI arguments for claude -p.
func buildClaudeArgs(cfg config.ClaudeCodeCfg, req coretypes.CompletionRequest) []string {
	args := []string{"--print", "--output-format", "json"}

	if req.SystemPrompt != "" {
		args = append(args, "--system-prompt", req.SystemPrompt)
	}

	if cfg.DisableTools {
		args = append(args, "--tools", "")
	} else {
		if len(cfg.AllowedTools) > 0 {
			args = append(args, "--allowedTools")
			args = append(args, cfg.AllowedTools...)
		}

		if len(cfg.DisallowedTools) > 0 {
			args = append(args, "--disallowedTools")
			args = append(args, cfg.DisallowedTools...)
		}
	}

	if cfg.PermissionMode != "" {
		args = append(args, "--permission-mode", cfg.PermissionMode)
	}

	if cfg.Model != "" {
		args = append(args, "--model", cfg.Model)
	}

	if cfg.FallbackModel != "" {
		args = append(args, "--fallback-model", cfg.FallbackModel)
	}

	if cfg.SessionID != "" {
		args = append(args, "--session-id", cfg.SessionID)
	}

	for _, dir := range cfg.AddDirs {
		args = append(args, "--add-dir", dir)
	}

	return args
}

// flattenMessages converts a conversation history into a single text prompt for stdin.
func flattenMessages(msgs []coretypes.Message) string {
	var b strings.Builder
	for _, m := range msgs {
		switch m.Role {
		case coretypes.RoleUser:
			fmt.Fprintf(&b, "User: %s\n\n", m.Content)
		case coretypes.RoleAssistant:
			fmt.Fprintf(&b, "Assistant: %s\n\n", m.Content)
		case coretypes.RoleTool:
			fmt.Fprintf(&b, "Tool result: %s\n\n", m.Content)
		}
	}
	return b.String()
}

// parseClaudeOutput parses the JSON output from `claude -p --output-format json`.
func parseClaudeOutput(
	stdout, stderr []byte,
	execErr error,
	elapsed time.Duration,
	log *slog.Logger,
) (coretypes.CompletionResponse, error) {
	// If the process failed and there's no stdout, report the error
	if execErr != nil && len(stdout) == 0 {
		errMsg := string(stderr)
		if errMsg == "" {
			errMsg = execErr.Error()
		}
		return coretypes.CompletionResponse{}, fmt.Errorf("claude-code process failed: %s", errMsg)
	}

	// Parse JSON output
	var output claudeJSONOutput
	if err := json.Unmarshal(stdout, &output); err != nil {
		// Fall back to treating stdout as plain text
		log.Warn("claude_code_json_parse_failed", "err", err, "stdout_len", len(stdout))
		return coretypes.CompletionResponse{
			Content:      strings.TrimSpace(string(stdout)),
			FinishReason: "stop",
		}, nil
	}

	if output.IsError {
		return coretypes.CompletionResponse{}, fmt.Errorf("claude-code error: %s", output.Result)
	}

	// Extract text from result field or content blocks
	content := output.Result
	if content == "" && len(output.ContentBlock) > 0 {
		var parts []string
		for _, block := range output.ContentBlock {
			if block.Type == "text" && block.Text != "" {
				parts = append(parts, block.Text)
			}
		}
		content = strings.Join(parts, "\n")
	}

	finishReason := "stop"
	if execErr != nil {
		finishReason = "error"
	}

	log.Info("claude_code_response",
		"content_len", len(content),
		"input_tokens", output.Usage.InputTokens,
		"output_tokens", output.Usage.OutputTokens,
		"num_turns", output.NumTurns,
		"cost_usd", output.TotalCost,
		"elapsed_ms", elapsed.Milliseconds(),
	)

	return coretypes.CompletionResponse{
		Content: content,
		Usage: coretypes.TokenUsage{
			InputTokens:  output.Usage.InputTokens,
			OutputTokens: output.Usage.OutputTokens,
			TotalTokens:  output.Usage.InputTokens + output.Usage.OutputTokens,
		},
		FinishReason: finishReason,
	}, nil
}

// filterEnv returns a copy of environ with the named keys removed.
func filterEnv(environ []string, keys ...string) []string {
	out := make([]string, 0, len(environ))
	for _, e := range environ {
		skip := false
		for _, k := range keys {
			if strings.HasPrefix(e, k+"=") {
				skip = true
				break
			}
		}
		if !skip {
			out = append(out, e)
		}
	}
	return out
}