fn_registry/functions/core/claude_stream.go

package core

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
)

// ClaudeEventType es el tipo discriminador de eventos del stream-json de claude -p.
type ClaudeEventType string

const (
	ClaudeEventSystem     ClaudeEventType = "system"
	ClaudeEventAssistant  ClaudeEventType = "assistant"
	ClaudeEventUser       ClaudeEventType = "user"        // tool_result
	ClaudeEventResult     ClaudeEventType = "result"      // final
	ClaudeEventToolUse    ClaudeEventType = "tool_use"    // sintetico
	ClaudeEventToolResult ClaudeEventType = "tool_result" // sintetico
	ClaudeEventTextDelta  ClaudeEventType = "text_delta"  // sintetico (porcion legible)
	ClaudeEventError      ClaudeEventType = "error"       // sintetico
)

// ClaudeEvent es un evento decodificado del stream. Raw siempre contiene la
// linea NDJSON original para casos no contemplados. Para los tipos comunes,
// los campos especificos vienen rellenos.
type ClaudeEvent struct {
	Type ClaudeEventType `json:"type"`
	Raw  json.RawMessage `json:"raw,omitempty"`

	// Para system/init
	Subtype   string `json:"subtype,omitempty"`
	SessionID string `json:"session_id,omitempty"`
	Model     string `json:"model,omitempty"`

	// Para text_delta (sintetico): porcion textual del mensaje del asistente
	Text string `json:"text,omitempty"`

	// Para tool_use (sintetico)
	ToolUseID string          `json:"tool_use_id,omitempty"`
	ToolName  string          `json:"tool_name,omitempty"`
	ToolInput json.RawMessage `json:"tool_input,omitempty"`

	// Para tool_result (sintetico)
	ToolResultID      string `json:"tool_result_id,omitempty"`
	ToolResultContent string `json:"tool_result_content,omitempty"`
	ToolResultIsError bool   `json:"tool_result_is_error,omitempty"`

	// Para result (final)
	StopReason string `json:"stop_reason,omitempty"`
	IsError    bool   `json:"is_error,omitempty"`
	Result     string `json:"result,omitempty"`

	// Para error
	Error string `json:"error,omitempty"`
}

// ClaudeStreamOpts configura el lanzamiento.
type ClaudeStreamOpts struct {
	Bin     string            // default "claude" si vacio
	Args    []string          // args extra (NO incluyas -p ni --output-format ni --verbose; se añaden automaticamente)
	Stdin   io.Reader         // prompt user (puede ser nil si Args lleva el prompt en posicional)
	Workdir string            // CWD del subprocess
	Env     map[string]string // env extra (se mergea con os.Environ())
	Stderr  io.Writer         // si != nil, recibe stderr del subprocess en vivo
}

// streamRawLine es la estructura minima para detectar el tipo de una linea NDJSON.
type streamRawLine struct {
	Type    ClaudeEventType `json:"type"`
	Subtype string          `json:"subtype,omitempty"`

	// system
	SessionID string `json:"session_id,omitempty"`
	Model     string `json:"model,omitempty"`

	// result
	StopReason string `json:"stop_reason,omitempty"`
	IsError    bool   `json:"is_error,omitempty"`
	Result     string `json:"result,omitempty"`

	// assistant / user
	Message *streamMessage `json:"message,omitempty"`
}

type streamMessage struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

type contentBlock 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"`
	IsError   bool            `json:"is_error,omitempty"`
}

// extractToolResultContent extrae el texto de un tool_result content que puede
// ser string o array de bloques [{type:text,text:"..."}].
func extractToolResultContent(raw json.RawMessage) string {
	if len(raw) == 0 {
		return ""
	}
	// Intentar como string
	var s string
	if err := json.Unmarshal(raw, &s); err == nil {
		return s
	}
	// Intentar como array de bloques
	var blocks []contentBlock
	if err := json.Unmarshal(raw, &blocks); err != nil {
		return string(raw)
	}
	var sb strings.Builder
	for _, b := range blocks {
		if b.Type == "text" {
			sb.WriteString(b.Text)
		}
	}
	return sb.String()
}

// StreamClaude lanza `claude -p --output-format stream-json --verbose <args...>`
// y retorna un canal de eventos. El canal se cierra cuando termina el subprocess
// (EOF en stdout). El cancel del ctx mata al subprocess (SIGTERM, luego SIGKILL).
//
// La goroutine interna se encarga de:
//   - Leer stdout linea a linea (NDJSON, buffer 4MB).
//   - Decodificar cada linea a un evento del protocolo claude.
//   - Para mensajes "assistant" expandir el array message.content emitiendo
//     ClaudeEventTextDelta por cada bloque text y ClaudeEventToolUse por cada
//     bloque tool_use.
//   - Para mensajes "user" detectar tool_result y emitir ClaudeEventToolResult.
//   - Si stdout emite linea no-JSON, emite ClaudeEventError con el contenido.
//   - Capturar el exit code del subprocess; si != 0 emite ClaudeEventError final.
//
// Retorna error si el spawn falla. Si retorna chan != nil, el caller DEBE leerlo
// hasta que se cierre o cancelar el ctx.
func StreamClaude(ctx context.Context, opts ClaudeStreamOpts) (<-chan ClaudeEvent, error) {
	bin := opts.Bin
	if bin == "" {
		var err error
		bin, err = exec.LookPath("claude")
		if err != nil {
			return nil, fmt.Errorf("claude binary not found: %w", err)
		}
	}

	args := append([]string{"-p", "--output-format", "stream-json", "--verbose"}, opts.Args...)
	cmd := exec.CommandContext(ctx, bin, args...)

	if opts.Stdin != nil {
		cmd.Stdin = opts.Stdin
	}

	if opts.Workdir != "" {
		cmd.Dir = opts.Workdir
	}

	// Merge env
	if len(opts.Env) > 0 {
		env := os.Environ()
		for k, v := range opts.Env {
			env = append(env, k+"="+v)
		}
		cmd.Env = env
	}

	// Stderr
	stderrWriter := io.Discard
	if opts.Stderr != nil {
		stderrWriter = opts.Stderr
	}

	// Capturar stderr para reportar en error final
	var stderrBuf strings.Builder
	cmd.Stderr = io.MultiWriter(stderrWriter, &stderrBuf)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start claude: %w", err)
	}

	ch := make(chan ClaudeEvent, 64)

	go func() {
		defer close(ch)

		scanner := bufio.NewScanner(stdout)
		scanner.Buffer(make([]byte, 4*1024*1024), 4*1024*1024)

		send := func(ev ClaudeEvent) {
			select {
			case ch <- ev:
			case <-ctx.Done():
			}
		}

		for scanner.Scan() {
			line := scanner.Bytes()
			if len(line) == 0 {
				continue
			}

			var raw streamRawLine
			if err := json.Unmarshal(line, &raw); err != nil {
				send(ClaudeEvent{
					Type:  ClaudeEventError,
					Error: fmt.Sprintf("non-json line: %s", string(line)),
					Raw:   json.RawMessage(line),
				})
				continue
			}

			switch raw.Type {
			case ClaudeEventSystem:
				send(ClaudeEvent{
					Type:      ClaudeEventSystem,
					Subtype:   raw.Subtype,
					SessionID: raw.SessionID,
					Model:     raw.Model,
					Raw:       json.RawMessage(line),
				})

			case ClaudeEventAssistant:
				if raw.Message != nil && len(raw.Message.Content) > 0 {
					var blocks []contentBlock
					if err := json.Unmarshal(raw.Message.Content, &blocks); err == nil {
						for _, b := range blocks {
							switch b.Type {
							case "text":
								send(ClaudeEvent{
									Type: ClaudeEventTextDelta,
									Text: b.Text,
									Raw:  json.RawMessage(line),
								})
							case "tool_use":
								send(ClaudeEvent{
									Type:      ClaudeEventToolUse,
									ToolUseID: b.ID,
									ToolName:  b.Name,
									ToolInput: b.Input,
									Raw:       json.RawMessage(line),
								})
							}
						}
					}
				}
				// Emitir tambien el evento assistant crudo
				send(ClaudeEvent{
					Type: ClaudeEventAssistant,
					Raw:  json.RawMessage(line),
				})

			case ClaudeEventUser:
				if raw.Message != nil && len(raw.Message.Content) > 0 {
					var blocks []contentBlock
					if err := json.Unmarshal(raw.Message.Content, &blocks); err == nil {
						for _, b := range blocks {
							if b.Type == "tool_result" {
								content := extractToolResultContent(b.Content)
								send(ClaudeEvent{
									Type:              ClaudeEventToolResult,
									ToolResultID:      b.ToolUseID,
									ToolResultContent: content,
									ToolResultIsError:  b.IsError,
									Raw:               json.RawMessage(line),
								})
							}
						}
					}
				}
				send(ClaudeEvent{
					Type: ClaudeEventUser,
					Raw:  json.RawMessage(line),
				})

			case ClaudeEventResult:
				send(ClaudeEvent{
					Type:       ClaudeEventResult,
					Subtype:    raw.Subtype,
					SessionID:  raw.SessionID,
					StopReason: raw.StopReason,
					IsError:    raw.IsError,
					Result:     raw.Result,
					Raw:        json.RawMessage(line),
				})

			default:
				send(ClaudeEvent{
					Type: raw.Type,
					Raw:  json.RawMessage(line),
				})
			}
		}

		// Esperar a que termine el subprocess
		if err := cmd.Wait(); err != nil {
			if ctx.Err() != nil {
				// Cancelado por contexto — salida limpia
				return
			}
			stderr := strings.TrimSpace(stderrBuf.String())
			errMsg := fmt.Sprintf("claude exit error: %v", err)
			if stderr != "" {
				// Solo las ultimas lineas para no saturar
				lines := strings.Split(stderr, "\n")
				tail := lines
				if len(lines) > 5 {
					tail = lines[len(lines)-5:]
				}
				errMsg = fmt.Sprintf("claude exit error: %v: %s", err, strings.Join(tail, "; "))
			}
			send(ClaudeEvent{
				Type:  ClaudeEventError,
				Error: errMsg,
			})
		}
	}()

	return ch, nil
}