summaryrefslogtreecommitdiff
path: root/internal/chatrun
diff options
context:
space:
mode:
authorPaul Buetow <paul@buetow.org>2026-06-11 00:28:27 +0300
committerPaul Buetow <paul@buetow.org>2026-06-11 00:28:27 +0300
commitb4f06c656d8d8733a9bf2e2e2a5eb2a7a48389bb (patch)
tree7df0eac26ded2b4854225711dfde8f40d80a0400 /internal/chatrun
parent2b2f4110da53a03742faff89cdec17c55e091a90 (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.go114
-rw-r--r--internal/chatrun/chatrun_test.go135
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"))
+ }
+}