Support prompt input streaming

Author: StephanSchmidt

pkg/claude/claude.go

@@ -29,6 +29,16 @@ const (
 	StreamJSONOutput OutputFormat = "stream-json"
 )
 
+// InputFormat defines the input format for Claude Code requests
+type InputFormat string
+
+const (
+	// TextInput sends plain text input (default)
+	TextInput InputFormat = "text"
+	// StreamJSONInput sends streaming JSON input for multiple prompts
+	StreamJSONInput InputFormat = "stream-json"
+)
+
 // ClaudeClient is the main client for interacting with Claude Code
 type ClaudeClient struct {
 	// BinPath is the path to the Claude Code binary
@@ -41,6 +51,8 @@ type ClaudeClient struct {
 type RunOptions struct {
 	// Format specifies the output format (text, json, stream-json)
 	Format OutputFormat
+	// InputFormat specifies the input format (text, stream-json)
+	InputFormat InputFormat
 	// SystemPrompt overrides the default system prompt
 	SystemPrompt string
 	// AppendPrompt appends to the default system prompt
@@ -126,6 +138,18 @@ type QueryOptions struct {
 	BufferConfig *buffer.Config  `json:"-"`
 }
 
+// StreamInputMessage represents a message for streaming input to Claude Code
+type StreamInputMessage struct {
+	Type    string              `json:"type"`
+	Message StreamInputContent  `json:"message"`
+}
+
+// StreamInputContent represents the content of a streaming input message
+type StreamInputContent struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
 // Message represents a message from Claude Code in streaming mode
 // Aligned with Python SDK message structure
 type Message struct {
@@ -321,8 +345,23 @@ func (c *ClaudeClient) RunPromptCtx(ctx context.Context, prompt string, opts *Ru
 	}, nil
 }
 
-// StreamPrompt executes a prompt with Claude Code and streams the results through a channel
+// StreamPrompt executes a prompt with Claude Code and streams the results thro	ugh a channel
 func (c *ClaudeClient) StreamPrompt(ctx context.Context, prompt string, opts *RunOptions) (<-chan Message, <-chan error) {
+	// Create a channel for a single prompt
+	promptCh := make(chan string, 1)
+	promptCh <- prompt
+	close(promptCh)
+	
+	// Use the multi-prompt streaming function
+	return c.StreamPromptsToSession(ctx, promptCh, opts)
+}
+
+// StreamPromptsToSession starts a Claude Code session and streams prompts to it continuously
+func (c *ClaudeClient) StreamPromptsToSession(
+	ctx context.Context,
+	promptCh <-chan string,
+	opts *RunOptions,
+) (<-chan Message, <-chan error) {
 	messageCh := make(chan Message)
 	errCh := make(chan error, 1)
 
@@ -333,11 +372,13 @@ func (c *ClaudeClient) StreamPrompt(ctx context.Context, prompt string, opts *Ru
 	// Force stream-json format for streaming
 	streamOpts := *opts
 	streamOpts.Format = StreamJSONOutput
+	streamOpts.InputFormat = StreamJSONInput
 
 	// Claude CLI requires --verbose when using --output-format=stream-json with --print
 	streamOpts.Verbose = true
 
-	args := BuildArgs(prompt, &streamOpts)
+	// Remove the initial prompt since we'll stream through stdin
+	args := BuildArgs("", &streamOpts)
 
 	go func() {
 		defer close(messageCh)
@@ -346,6 +387,12 @@ func (c *ClaudeClient) StreamPrompt(ctx context.Context, prompt string, opts *Ru
 		// Create a custom command that supports context
 		cmd := execCommand(ctx, c.BinPath, args...)
 
+		stdin, err := cmd.StdinPipe()
+		if err != nil {
+			errCh <- fmt.Errorf("failed to get stdin pipe: %w", err)
+			return
+		}
+
 		stdout, err := cmd.StdoutPipe()
 		if err != nil {
 			errCh <- fmt.Errorf("failed to get stdout pipe: %w", err)
@@ -376,6 +423,67 @@ func (c *ClaudeClient) StreamPrompt(ctx context.Context, prompt string, opts *Ru
 			return
 		}
 
+		// Channel to signal when command is ready to receive input
+		cmdReady := make(chan struct{})
+		
+		// Start a goroutine to handle input prompts
+		go func() {
+			defer stdin.Close()
+			
+			// Wait for command to be ready before processing prompts
+			select {
+			case <-cmdReady:
+				// Command is ready, proceed with prompt processing
+			case <-ctx.Done():
+				return
+			}
+			
+			for {
+				select {
+				case prompt, ok := <-promptCh:
+					if !ok {
+						// Prompt channel closed, close stdin
+						return
+					}
+					
+					// Create JSON message for streaming input
+					streamMsg := StreamInputMessage{
+						Type: "user",
+						Message: StreamInputContent{
+							Role:    "user",
+							Content: prompt,
+						},
+					}
+					
+					// Encode as JSON and send
+					jsonData, err := json.Marshal(streamMsg)
+					if err != nil {
+						errCh <- fmt.Errorf("failed to marshal stream input message: %w", err)
+						return
+					}
+					
+					if _, err := fmt.Fprintln(stdin, string(jsonData)); err != nil {
+						errCh <- fmt.Errorf("failed to write JSON prompt to stdin: %w", err)
+						return
+					}
+				case <-ctx.Done():
+					return
+				}
+			}
+		}()
+		
+		// Give the command a moment to initialize before signaling ready
+		// This prevents the race condition where we write to stdin before 
+		// the Claude process is ready to read from it
+		go func() {
+			select {
+			case <-time.After(100 * time.Millisecond):
+				close(cmdReady)
+			case <-ctx.Done():
+				return
+			}
+		}()
+
 		// Use buffered reader with configurable buffer size instead of scanner
 		reader := bufio.NewReaderSize(stdout, int(bufferConfig.MaxStdoutSize/1000)) // Use reasonable buffer size
 
@@ -642,6 +750,10 @@ func BuildArgs(prompt string, opts *RunOptions) []string {
 		args = append(args, "--output-format", string(opts.Format))
 	}
 
+	if opts.InputFormat != "" {
+		args = append(args, "--input-format", string(opts.InputFormat))
+	}
+
 	if opts.SystemPrompt != "" {
 		args = append(args, "--system-prompt", opts.SystemPrompt)
 	}

pkg/claude/input_format_test.go

@@ -0,0 +1,84 @@
+package claude
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestInputFormat_Constants(t *testing.T) {
+	tests := []struct {
+		format   InputFormat
+		expected string
+	}{
+		{TextInput, "text"},
+		{StreamJSONInput, "stream-json"},
+	}
+
+	for _, test := range tests {
+		if string(test.format) != test.expected {
+			t.Errorf("Expected %s, got %s", test.expected, string(test.format))
+		}
+	}
+}
+
+func TestBuildArgs_InputFormat(t *testing.T) {
+	tests := []struct {
+		name        string
+		opts        *RunOptions
+		expectedArg string
+	}{
+		{
+			name: "No input format specified",
+			opts: &RunOptions{},
+			expectedArg: "",
+		},
+		{
+			name: "Text input format",
+			opts: &RunOptions{InputFormat: TextInput},
+			expectedArg: "--input-format text",
+		},
+		{
+			name: "Stream JSON input format", 
+			opts: &RunOptions{InputFormat: StreamJSONInput},
+			expectedArg: "--input-format stream-json",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			args := BuildArgs("test prompt", test.opts)
+			argsStr := strings.Join(args, " ")
+			
+			if test.expectedArg == "" {
+				if strings.Contains(argsStr, "--input-format") {
+					t.Errorf("Expected no --input-format flag, but found one in: %s", argsStr)
+				}
+			} else {
+				if !strings.Contains(argsStr, test.expectedArg) {
+					t.Errorf("Expected to find '%s' in args: %s", test.expectedArg, argsStr)
+				}
+			}
+		})
+	}
+}
+
+func TestStreamInputMessage_JSONMarshaling(t *testing.T) {
+	msg := StreamInputMessage{
+		Type: "user",
+		Message: StreamInputContent{
+			Role:    "user",
+			Content: "Test message",
+		},
+	}
+
+	// Test the struct fields are correctly defined
+	if msg.Type != "user" {
+		t.Errorf("Expected Type 'user', got %s", msg.Type)
+	}
+	if msg.Message.Role != "user" {
+		t.Errorf("Expected Role 'user', got %s", msg.Message.Role)
+	}
+	if msg.Message.Content != "Test message" {
+		t.Errorf("Expected Content 'Test message', got %s", msg.Message.Content)
+	}
+}

pkg/claude/query_test.go

@@ -3,6 +3,9 @@ package claude
 import (
 	"context"
 	"encoding/json"
+	"os"
+	"os/exec"
+	"path/filepath"
 	"testing"
 	"time"
 
@@ -13,8 +16,9 @@ func TestQuery_PythonSDKAlignment(t *testing.T) {
 	if testing.Short() {
 		t.Skip("Skipping integration test in short mode")
 	}
-
-	client := &ClaudeClient{BinPath: "echo"}
+	
+	skipIfNoClaudeCLI(t)
+	client := newTestClient(t)
 	ctx := context.Background()
 
 	// Test basic Query method with QueryOptions
@@ -50,7 +54,8 @@ func TestQuerySync_PythonSDKAlignment(t *testing.T) {
 		t.Skip("Skipping integration test in short mode")
 	}
 
-	client := &ClaudeClient{BinPath: "echo"}
+	skipIfNoClaudeCLI(t)
+	client := newTestClient(t)
 	ctx := context.Background()
 
 	// Test synchronous Query method
@@ -198,4 +203,68 @@ func contains(slice []string, item string) bool {
 		}
 	}
 	return false
+}
+
+// Local test helper functions to avoid import cycle with test/utils
+
+// skipIfNoClaudeCLI skips the test if Claude Code CLI is not available
+func skipIfNoClaudeCLI(t *testing.T) {
+	if os.Getenv("USE_MOCK_SERVER") == "1" {
+		return // Mock server tests don't need Claude CLI
+	}
+	
+	claudePath := getTestClaudePath(t)
+	if _, err := exec.LookPath(claudePath); err != nil {
+		t.Skipf("Skipping test: Claude Code CLI not found at '%s'. Install Claude Code CLI or use mock server with USE_MOCK_SERVER=1", claudePath)
+	}
+	
+	// Test if Claude CLI is working (it will handle auth automatically)
+	cmd := exec.Command(claudePath, "--help")
+	if err := cmd.Run(); err != nil {
+		t.Skip("Skipping test: Claude Code CLI not working. Please ensure it's properly installed.")
+	}
+}
+
+// newTestClient creates a Claude client for testing
+func newTestClient(t *testing.T) *ClaudeClient {
+	return NewClient(getTestClaudePath(t))
+}
+
+// getTestClaudePath returns the path to Claude CLI for testing
+func getTestClaudePath(t *testing.T) string {
+	if path := os.Getenv("CLAUDE_CODE_PATH"); path != "" {
+		return path
+	}
+	
+	// Check if mock server is preferred
+	if os.Getenv("USE_MOCK_SERVER") == "1" {
+		return createMockClaudeScript(t)
+	}
+	
+	return "claude" // Default Claude Code CLI
+}
+
+// createMockClaudeScript creates a shell script that returns valid JSON
+func createMockClaudeScript(t *testing.T) string {
+	tempDir := t.TempDir()
+	mockScript := filepath.Join(tempDir, "mock-claude.sh")
+	
+	content := `#!/bin/bash
+# Mock Claude CLI that returns appropriate JSON response based on format
+if [[ "$*" == *"--format json"* ]]; then
+  # For regular JSON format (QuerySync)
+  echo '{"result": "Mock response", "sessionId": "test-session", "costUSD": 0.001, "durationMS": 100}'
+else
+  # For streaming format or default (Query)
+  echo '{"type": "system", "subtype": "init", "content": "Initializing"}'
+  echo '{"type": "assistant", "content": "Mock streaming response"}'
+  echo '{"type": "result", "result": "Mock response", "sessionId": "test-session", "costUSD": 0.001, "durationMS": 100}'
+fi
+`
+	
+	if err := os.WriteFile(mockScript, []byte(content), 0755); err != nil {
+		t.Fatalf("Failed to create mock script: %v", err)
+	}
+	
+	return mockScript
 }