diff options
| author | Paul Buetow <paul@buetow.org> | 2026-06-11 00:28:27 +0300 |
|---|---|---|
| committer | Paul Buetow <paul@buetow.org> | 2026-06-11 00:28:27 +0300 |
| commit | b4f06c656d8d8733a9bf2e2e2a5eb2a7a48389bb (patch) | |
| tree | 7df0eac26ded2b4854225711dfde8f40d80a0400 /internal/chatrun | |
| parent | 2b2f4110da53a03742faff89cdec17c55e091a90 (diff) | |
Extract shared chatrun package to eliminate chat-runner DRY violations
The CLI, LSP server and tmux code-action tool each carried near-identical
chat-running logic: invoking the LLM (streaming-aware), collecting the
response, and accounting sent/received bytes into the stats package.
Introduce internal/chatrun with:
- Invoke: streaming-aware LLM call that collects the full response and
optionally mirrors chunks to a writer (nil writer = collect only).
- SentBytes / Account: shared byte counting and stats.Update.
Wire all three surfaces to it:
- hexaicli: runChatRequest delegates to chatrun.Invoke; summarizeChatRun
uses chatrun.Account. Removed the duplicated streaming/simple helpers.
- hexaiaction: runOnce uses chatrun.Invoke + chatrun.Account, keeping the
tmux status update local.
- lsp: chatWithStats and the completion path use chatrun.SentBytes/Invoke;
extracted unavailableClientError to keep chatWithStats small.
chatrun has 100% test coverage; full suite passes with -race.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Diffstat (limited to 'internal/chatrun')
| -rw-r--r-- | internal/chatrun/chatrun.go | 114 | ||||
| -rw-r--r-- | internal/chatrun/chatrun_test.go | 135 |
2 files changed, 249 insertions, 0 deletions
diff --git a/internal/chatrun/chatrun.go b/internal/chatrun/chatrun.go new file mode 100644 index 0000000..01400cd --- /dev/null +++ b/internal/chatrun/chatrun.go @@ -0,0 +1,114 @@ +// Package chatrun holds the shared chat-running primitives that used to be +// duplicated across the three Hexai surfaces: the CLI (internal/hexaicli), the +// LSP server (internal/lsp) and the tmux code-action tool (internal/hexaiaction). +// +// Every surface had to do the same two things when talking to an LLM: +// +// 1. Invoke the client, preferring streaming (llm.Streamer) when the provider +// implements it and falling back to a single Chat call otherwise, while +// collecting the full response text. +// 2. Account for the exchange: count the bytes sent (sum of message contents) +// and received (response length) and feed them into the stats package. +// +// Both pieces lived in slightly diverging copies (runChatRequest/runStreaming +// Chat/runSimpleChat in the CLI, chatWithStats in the LSP, runOnce in the +// action tool). This package centralises them so all three share one +// implementation and behaviour stays identical. +package chatrun + +import ( + "context" + "fmt" + "io" + "strings" + + "codeberg.org/snonux/hexai/internal/llm" + "codeberg.org/snonux/hexai/internal/stats" +) + +// Chatter is the minimal client capability Invoke needs: a single Chat call. +// It is deliberately narrower than llm.Client (no Name/DefaultModel) so callers +// that only have a chat-doer abstraction — such as the code-action surface — +// can use Invoke without widening their own interface. Clients that also +// implement llm.Streamer get streaming automatically via a type assertion. +type Chatter interface { + Chat(ctx context.Context, messages []llm.Message, opts ...llm.RequestOption) (string, error) +} + +// Invoke sends msgs to the client and returns the full assistant response. +// +// When the client implements llm.Streamer the response is streamed and, if out +// is non-nil, each chunk is forwarded to out as it arrives (this is how the CLI +// renders incremental output). Otherwise a single Chat call is made and, when +// out is non-nil, the whole response is written to it once. +// +// Passing a nil out collects the response without writing it anywhere, which is +// what the LSP and code-action surfaces need (they post-process the text before +// applying it to a document). +func Invoke(ctx context.Context, client Chatter, msgs []llm.Message, opts []llm.RequestOption, out io.Writer) (string, error) { + if streamer, ok := client.(llm.Streamer); ok { + return invokeStreaming(ctx, streamer, msgs, opts, out) + } + return invokeSimple(ctx, client, msgs, opts, out) +} + +// invokeStreaming drives ChatStream, accumulating the full text while +// optionally mirroring each chunk to out. A write error to out aborts further +// writes and is returned once streaming finishes. +func invokeStreaming(ctx context.Context, client llm.Streamer, msgs []llm.Message, opts []llm.RequestOption, out io.Writer) (string, error) { + var output strings.Builder + var writeErr error + err := client.ChatStream(ctx, msgs, func(chunk string) { + output.WriteString(chunk) + if out == nil || writeErr != nil { + return + } + if _, werr := io.WriteString(out, chunk); werr != nil { + writeErr = werr + } + }, opts...) + if err != nil { + return "", err + } + if writeErr != nil { + return "", writeErr + } + return output.String(), nil +} + +// invokeSimple performs a single Chat call and, when out is non-nil, writes the +// whole response to it. +func invokeSimple(ctx context.Context, client Chatter, msgs []llm.Message, opts []llm.RequestOption, out io.Writer) (string, error) { + output, err := client.Chat(ctx, msgs, opts...) + if err != nil { + return "", err + } + if out != nil { + if _, werr := fmt.Fprint(out, output); werr != nil { + return "", werr + } + } + return output, nil +} + +// SentBytes returns the total number of content bytes across msgs. This is the +// "sent" figure every surface reports and feeds into the stats package. +func SentBytes(msgs []llm.Message) int { + sent := 0 + for _, m := range msgs { + sent += len(m.Content) + } + return sent +} + +// Account records a completed exchange in the stats package and returns the +// sent/received byte counts so callers can build their own summaries. The +// stats.Update error is intentionally swallowed because none of the surfaces +// treat a stats failure as fatal; callers that want to log it can call +// stats.Update directly instead. +func Account(ctx context.Context, provider, model string, msgs []llm.Message, output string) (sent, recv int) { + sent = SentBytes(msgs) + recv = len(output) + _ = stats.Update(ctx, provider, model, sent, recv) + return sent, recv +} diff --git a/internal/chatrun/chatrun_test.go b/internal/chatrun/chatrun_test.go new file mode 100644 index 0000000..f2f5806 --- /dev/null +++ b/internal/chatrun/chatrun_test.go @@ -0,0 +1,135 @@ +package chatrun + +import ( + "bytes" + "context" + "errors" + "strings" + "testing" + + "codeberg.org/snonux/hexai/internal/llm" +) + +// simpleClient implements only Chatter (no streaming). +type simpleClient struct { + resp string + err error +} + +func (c simpleClient) Chat(_ context.Context, _ []llm.Message, _ ...llm.RequestOption) (string, error) { + return c.resp, c.err +} + +// streamClient implements Chatter and llm.Streamer. +type streamClient struct { + chunks []string + err error +} + +func (c streamClient) Chat(_ context.Context, _ []llm.Message, _ ...llm.RequestOption) (string, error) { + return strings.Join(c.chunks, ""), nil +} + +func (c streamClient) ChatStream(_ context.Context, _ []llm.Message, onDelta func(string), _ ...llm.RequestOption) error { + if c.err != nil { + return c.err + } + for _, chunk := range c.chunks { + onDelta(chunk) + } + return nil +} + +// errWriter always fails, to exercise the write-error paths. +type errWriter struct{ err error } + +func (w errWriter) Write([]byte) (int, error) { return 0, w.err } + +func TestInvoke_SimpleCollectsAndWrites(t *testing.T) { + var out bytes.Buffer + got, err := Invoke(context.Background(), simpleClient{resp: "hello"}, nil, nil, &out) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != "hello" || out.String() != "hello" { + t.Fatalf("got=%q out=%q, want hello/hello", got, out.String()) + } +} + +func TestInvoke_SimpleNilWriter(t *testing.T) { + got, err := Invoke(context.Background(), simpleClient{resp: "x"}, nil, nil, nil) + if err != nil || got != "x" { + t.Fatalf("got=%q err=%v, want x/nil", got, err) + } +} + +func TestInvoke_SimpleChatError(t *testing.T) { + _, err := Invoke(context.Background(), simpleClient{err: errors.New("boom")}, nil, nil, nil) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("expected chat error, got %v", err) + } +} + +func TestInvoke_SimpleWriteError(t *testing.T) { + _, err := Invoke(context.Background(), simpleClient{resp: "x"}, nil, nil, errWriter{err: errors.New("wfail")}) + if err == nil || !strings.Contains(err.Error(), "wfail") { + t.Fatalf("expected write error, got %v", err) + } +} + +func TestInvoke_StreamingCollectsAndWrites(t *testing.T) { + var out bytes.Buffer + got, err := Invoke(context.Background(), streamClient{chunks: []string{"a", "b", "c"}}, nil, nil, &out) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != "abc" || out.String() != "abc" { + t.Fatalf("got=%q out=%q, want abc/abc", got, out.String()) + } +} + +func TestInvoke_StreamingNilWriter(t *testing.T) { + got, err := Invoke(context.Background(), streamClient{chunks: []string{"a", "b"}}, nil, nil, nil) + if err != nil || got != "ab" { + t.Fatalf("got=%q err=%v, want ab/nil", got, err) + } +} + +func TestInvoke_StreamingError(t *testing.T) { + _, err := Invoke(context.Background(), streamClient{err: errors.New("sfail")}, nil, nil, nil) + if err == nil || !strings.Contains(err.Error(), "sfail") { + t.Fatalf("expected stream error, got %v", err) + } +} + +func TestInvoke_StreamingWriteError(t *testing.T) { + // Full text must still be collected even though the writer fails. + _, err := Invoke(context.Background(), streamClient{chunks: []string{"a", "b"}}, nil, nil, errWriter{err: errors.New("wfail")}) + if err == nil || !strings.Contains(err.Error(), "wfail") { + t.Fatalf("expected write error, got %v", err) + } +} + +func TestSentBytes(t *testing.T) { + msgs := []llm.Message{ + {Role: "system", Content: "abc"}, + {Role: "user", Content: "de"}, + } + if got := SentBytes(msgs); got != 5 { + t.Fatalf("SentBytes = %d, want 5", got) + } + if got := SentBytes(nil); got != 0 { + t.Fatalf("SentBytes(nil) = %d, want 0", got) + } +} + +func TestAccount(t *testing.T) { + msgs := []llm.Message{{Role: "user", Content: "1234"}} + sent, recv := Account(context.Background(), "prov", "model", msgs, "response") + if sent != 4 { + t.Fatalf("sent = %d, want 4", sent) + } + if recv != len("response") { + t.Fatalf("recv = %d, want %d", recv, len("response")) + } +} |
