From 02e79cb7b49b010ed66c61e488e9c087c8cfd993 Mon Sep 17 00:00:00 2001 From: Justin McCarthy Date: Mon, 16 Jun 2025 11:32:30 -0700 Subject: [PATCH] Implement remaining TODOs --- README.md | 10 +++++++++- claudecode/query.go | 2 +- examples/quick_start.go | 5 ++++- internal/client.go | 24 ++++++++++++++---------- internal/client_test.go | 30 ++++++++++++++++++++++++++++-- internal/transport.go | 22 ++++++++++++++++++++++ internal/transport_test.go | 11 ++++++++++- model/types.go | 1 + 8 files changed, 89 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 794dbfea9..41d74e5ba 100644 --- a/README.md +++ b/README.md @@ -29,13 +29,16 @@ import ( func main() { ctx := context.Background() - ch, err := claudecode.Query(ctx, "What is 2 + 2?", nil) + ch, errCh, err := claudecode.Query(ctx, "What is 2 + 2?", nil) if err != nil { log.Fatal(err) } for msg := range ch { log.Printf("%+v", msg) } + if e := <-errCh; e != nil { + log.Fatalf("query error: %v", e) + } } ``` @@ -43,6 +46,11 @@ func main() { See `examples/quick_start.go` for a more complete example of using the package. +### Custom CLI Path + +If the `claude` CLI binary is not on your `PATH`, set `Options.CLIPath` to the +location of the executable when calling `Query`. + ## License MIT diff --git a/claudecode/query.go b/claudecode/query.go index 51cb4496f..8f28769b7 100644 --- a/claudecode/query.go +++ b/claudecode/query.go @@ -9,7 +9,7 @@ import ( ) // Query sends a prompt to Claude Code and returns a stream of Messages. -func Query(ctx context.Context, prompt string, opts *model.Options) (<-chan model.Message, error) { +func Query(ctx context.Context, prompt string, opts *model.Options) (<-chan model.Message, <-chan error, error) { os.Setenv("CLAUDE_CODE_ENTRYPOINT", "sdk-go") transport := &internal.SubprocessCLITransport{} client := &internal.Client{Transport: transport} diff --git a/examples/quick_start.go b/examples/quick_start.go index 4fddad049..0d5c76631 100644 --- a/examples/quick_start.go +++ b/examples/quick_start.go @@ -10,7 +10,7 @@ import ( func main() { ctx := context.Background() - ch, err := claudecode.Query(ctx, "What is 2 + 2?", nil) + ch, errCh, err := claudecode.Query(ctx, "What is 2 + 2?", nil) if err != nil { log.Fatal(err) } @@ -18,4 +18,7 @@ func main() { for msg := range ch { log.Printf("%+v", msg) } + if e := <-errCh; e != nil { + log.Fatalf("query error: %v", e) + } } diff --git a/internal/client.go b/internal/client.go index 6c0db25a1..9ff100e56 100644 --- a/internal/client.go +++ b/internal/client.go @@ -12,35 +12,39 @@ type Client struct { } // Query connects to the transport and streams parsed messages. -func (c *Client) Query(ctx context.Context, prompt string, opts *model.Options) (<-chan model.Message, error) { - if sp, ok := c.Transport.(*SubprocessCLITransport); ok { - sp.Prompt = prompt - sp.Options = opts - if opts != nil { - sp.Cwd = opts.Cwd - } +func (c *Client) Query(ctx context.Context, prompt string, opts *model.Options) (<-chan model.Message, <-chan error, error) { + if err := c.Transport.SendRequest(ctx, prompt, opts); err != nil { + return nil, nil, err } if err := c.Transport.Connect(ctx); err != nil { - return nil, err + return nil, nil, err } rawCh, err := c.Transport.ReceiveMessages(ctx) if err != nil { - return nil, err + return nil, nil, err } out := make(chan model.Message) + errCh := make(chan error, 1) go func() { defer close(out) + defer close(errCh) defer c.Transport.Disconnect() for data := range rawCh { + if errVal, ok := data["error"]; ok { + if e, ok2 := errVal.(error); ok2 { + errCh <- e + return + } + } if m := parseMessage(data); m != nil { out <- m } } }() - return out, nil + return out, errCh, nil } func parseMessage(data map[string]any) model.Message { diff --git a/internal/client_test.go b/internal/client_test.go index fb9046dff..0668d1d4a 100644 --- a/internal/client_test.go +++ b/internal/client_test.go @@ -11,6 +11,7 @@ type stubTransport struct { msgs []map[string]any connected bool disconnected bool + sent bool } func (s *stubTransport) Connect(ctx context.Context) error { @@ -23,6 +24,11 @@ func (s *stubTransport) Disconnect() error { return nil } +func (s *stubTransport) SendRequest(ctx context.Context, prompt string, opts *model.Options) error { + s.sent = true + return nil +} + func (s *stubTransport) ReceiveMessages(ctx context.Context) (<-chan map[string]any, error) { ch := make(chan map[string]any) go func() { @@ -91,7 +97,7 @@ func TestClientQuery(t *testing.T) { {"type": "result", "subtype": "done"}, }} c := &Client{Transport: st} - ch, err := c.Query(context.Background(), "hi", nil) + ch, errCh, err := c.Query(context.Background(), "hi", nil) if err != nil { t.Fatal(err) } @@ -102,7 +108,27 @@ func TestClientQuery(t *testing.T) { if len(msgs) != 2 { t.Fatalf("expected 2 messages, got %d", len(msgs)) } - if !st.connected || !st.disconnected { + if !st.connected || !st.disconnected || !st.sent { t.Fatalf("transport lifecycle not called") } + if errVal := <-errCh; errVal != nil { + t.Fatalf("unexpected error: %v", errVal) + } +} + +func TestClientQueryErrorPropagation(t *testing.T) { + perr := &model.ProcessError{Msg: "fail", ExitCode: 1} + st := &stubTransport{msgs: []map[string]any{ + {"error": perr}, + }} + c := &Client{Transport: st} + ch, errCh, err := c.Query(context.Background(), "hi", nil) + if err != nil { + t.Fatal(err) + } + for range ch { + } + if e := <-errCh; e != perr { + t.Fatalf("expected error propagation") + } } diff --git a/internal/transport.go b/internal/transport.go index 9448dfcdc..d6feda877 100644 --- a/internal/transport.go +++ b/internal/transport.go @@ -19,6 +19,7 @@ import ( type Transport interface { Connect(ctx context.Context) error Disconnect() error + SendRequest(ctx context.Context, prompt string, opts *model.Options) error ReceiveMessages(ctx context.Context) (<-chan map[string]any, error) } @@ -34,6 +35,19 @@ type SubprocessCLITransport struct { stderr io.ReadCloser } +// SendRequest stores the prompt and options for the next connection. +func (t *SubprocessCLITransport) SendRequest(ctx context.Context, prompt string, opts *model.Options) error { + t.Prompt = prompt + t.Options = opts + if opts != nil { + t.Cwd = opts.Cwd + if opts.CLIPath != "" { + t.CLIPath = opts.CLIPath + } + } + return nil +} + // Connect starts the CLI process with the provided options. func (t *SubprocessCLITransport) Connect(ctx context.Context) error { if t.CLIPath == "" { @@ -153,6 +167,14 @@ func buildCommand(cliPath, prompt string, opts *model.Options) []string { if len(opts.AllowedTools) > 0 { args = append(args, "--allowedTools", join(opts.AllowedTools)) } + if opts.MaxThinkingTokens > 0 { + args = append(args, "--max-thinking-tokens", fmt.Sprint(opts.MaxThinkingTokens)) + } else { + args = append(args, "--max-thinking-tokens", "8000") + } + if len(opts.MCPTools) > 0 { + args = append(args, "--mcpTools", join(opts.MCPTools)) + } if opts.MaxTurns > 0 { args = append(args, "--max-turns", fmt.Sprint(opts.MaxTurns)) } diff --git a/internal/transport_test.go b/internal/transport_test.go index 727ef6da0..38b8f7fb1 100644 --- a/internal/transport_test.go +++ b/internal/transport_test.go @@ -10,7 +10,16 @@ import ( func TestBuildCommandBasic(t *testing.T) { opts := &model.Options{SystemPrompt: "hi"} cmd := buildCommand("/usr/bin/claude", "hello", opts) - expect := []string{"/usr/bin/claude", "--output-format", "stream-json", "--verbose", "--system-prompt", "hi", "--print", "hello"} + expect := []string{"/usr/bin/claude", "--output-format", "stream-json", "--verbose", "--system-prompt", "hi", "--max-thinking-tokens", "8000", "--print", "hello"} + if !reflect.DeepEqual(cmd, expect) { + t.Fatalf("unexpected cmd: %v", cmd) + } +} + +func TestBuildCommandWithExtras(t *testing.T) { + opts := &model.Options{MaxThinkingTokens: 9000, MCPTools: []string{"foo", "bar"}} + cmd := buildCommand("/usr/bin/claude", "hello", opts) + expect := []string{"/usr/bin/claude", "--output-format", "stream-json", "--verbose", "--max-thinking-tokens", "9000", "--mcpTools", "foo,bar", "--print", "hello"} if !reflect.DeepEqual(cmd, expect) { t.Fatalf("unexpected cmd: %v", cmd) } diff --git a/model/types.go b/model/types.go index 44050ab32..f161364fa 100644 --- a/model/types.go +++ b/model/types.go @@ -86,6 +86,7 @@ type Options struct { AppendSystemPrompt string MCPTools []string MCPServers map[string]MCPServerConfig + CLIPath string PermissionMode PermissionMode ContinueConversation bool Resume string