7 files changed, 301 insertions, 74 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"))
+	}
+}
diff --git a/internal/hexaiaction/prompts.go b/internal/hexaiaction/prompts.go
index 4b2d8be..c5e496b 100644
--- a/internal/hexaiaction/prompts.go
+++ b/internal/hexaiaction/prompts.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"codeberg.org/snonux/hexai/internal/appconfig"
+	"codeberg.org/snonux/hexai/internal/chatrun"
 	"codeberg.org/snonux/hexai/internal/llm"
 	"codeberg.org/snonux/hexai/internal/llmutils"
 	"codeberg.org/snonux/hexai/internal/stats"
@@ -115,25 +116,27 @@ func runCustom(ctx context.Context, cfg actionConfig, client chatDoer, ca appcon
 // runOnce sends a single system+user prompt pair to the LLM, strips code
 // fences from the response, records stats, and updates the tmux status line.
 // Pass a zero-value requestArgs{} when no extra options are needed.
+//
+// The LLM invocation and byte/stats accounting are delegated to the shared
+// chatrun package so this surface stays in lock-step with the CLI and LSP. The
+// tmux status update is kept here because it is specific to the action tool's
+// reporting style.
 func runOnce(ctx context.Context, client chatDoer, sys, user string, req requestArgs) (string, error) {
 	msgs := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: user}}
-	txt, err := client.Chat(ctx, msgs, req.options...)
+	// runOnce never streams to a writer, so pass nil out.
+	txt, err := chatrun.Invoke(ctx, client, msgs, req.options, nil)
 	if err != nil {
 		return "", err
 	}
 	out := strings.TrimSpace(StripFences(txt))
-	// Contribute to global stats and update tmux status
-	sent := 0
-	for _, m := range msgs {
-		sent += len(m.Content)
-	}
-	recv := len(out)
 	model := strings.TrimSpace(req.model)
 	if model == "" {
 		model = client.DefaultModel()
 	}
 	provider := providerOf(client)
-	_ = stats.Update(ctx, provider, model, sent, recv)
+	// Account for the exchange against the post-fence-strip response so the
+	// recorded recv bytes match what the user actually receives.
+	chatrun.Account(ctx, provider, model, msgs, out)
 	if snap, err := stats.TakeSnapshot(); err == nil {
 		scopeReqs := snap.ScopeReqs(provider, model)
 		scopeRPM := snap.ScopeRPM(provider, model)
diff --git a/internal/hexaicli/run.go b/internal/hexaicli/run.go
index 94dcd12..6614bc5 100644
--- a/internal/hexaicli/run.go
+++ b/internal/hexaicli/run.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"codeberg.org/snonux/hexai/internal/appconfig"
+	"codeberg.org/snonux/hexai/internal/chatrun"
 	"codeberg.org/snonux/hexai/internal/llm"
 	"codeberg.org/snonux/hexai/internal/llmutils"
 	"codeberg.org/snonux/hexai/internal/logging"
@@ -503,51 +504,16 @@ func effectiveModel(req requestArgs, client llm.Client) string {
 	return model
 }
 
+// runChatRequest invokes the LLM via the shared chatrun.Invoke helper, which
+// prefers streaming when the provider supports it and mirrors chunks to out.
 func runChatRequest(ctx context.Context, client llm.Client, req requestArgs, msgs []llm.Message, out io.Writer) (string, error) {
-	if streamer, ok := client.(llm.Streamer); ok {
-		return runStreamingChat(ctx, streamer, msgs, req.options, out)
-	}
-	return runSimpleChat(ctx, client, msgs, req.options, out)
-}
-
-func runStreamingChat(ctx context.Context, client llm.Streamer, msgs []llm.Message, options []llm.RequestOption, out io.Writer) (string, error) {
-	var output strings.Builder
-	var writeErr error
-	if err := client.ChatStream(ctx, msgs, func(chunk string) {
-		if writeErr != nil {
-			return
-		}
-		output.WriteString(chunk)
-		if _, err := fmt.Fprint(out, chunk); err != nil {
-			writeErr = err
-		}
-	}, options...); err != nil {
-		return "", err
-	}
-	if writeErr != nil {
-		return "", writeErr
-	}
-	return output.String(), nil
-}
-
-func runSimpleChat(ctx context.Context, client llm.Client, msgs []llm.Message, options []llm.RequestOption, out io.Writer) (string, error) {
-	output, err := client.Chat(ctx, msgs, options...)
-	if err != nil {
-		return "", err
-	}
-	if _, err := fmt.Fprint(out, output); err != nil {
-		return "", err
-	}
-	return output, nil
+	return chatrun.Invoke(ctx, client, msgs, req.options, out)
 }
 
 func summarizeChatRun(ctx context.Context, client llm.Client, model string, msgs []llm.Message, output string) chatRunSummary {
 	summary := chatRunSummary{snapshot: stats.Snapshot{Window: time.Hour}}
-	for _, m := range msgs {
-		summary.sent += len(m.Content)
-	}
-	summary.recv = len(output)
-	_ = stats.Update(ctx, client.Name(), model, summary.sent, summary.recv)
+	// Byte accounting + stats.Update are shared via chatrun.Account.
+	summary.sent, summary.recv = chatrun.Account(ctx, client.Name(), model, msgs, output)
 	snap, err := stats.TakeSnapshot()
 	if err == nil {
 		summary.snapshot = snap
diff --git a/internal/hexaicli/run_output_test.go b/internal/hexaicli/run_output_test.go
index e61e4b6..07a2cab 100644
--- a/internal/hexaicli/run_output_test.go
+++ b/internal/hexaicli/run_output_test.go
@@ -227,7 +227,7 @@ func TestRunStreamingChat_StreamError(t *testing.T) {
 		streamErr:  fmt.Errorf("stream broken"),
 	}
 	var out bytes.Buffer
-	_, err := runStreamingChat(context.Background(), client, nil, nil, &out)
+	_, err := runChatRequest(context.Background(), client, requestArgs{}, nil, &out)
 	if err == nil || !strings.Contains(err.Error(), "stream broken") {
 		t.Fatalf("expected stream error, got %v", err)
 	}
@@ -247,7 +247,7 @@ func (s *streamWriteErrClient) ChatStream(_ context.Context, _ []llm.Message, on
 func TestRunStreamingChat_WriteError(t *testing.T) {
 	client := &streamWriteErrClient{fakeClient: fakeClient{name: "p", model: "m"}}
 	w := errWriter{err: errors.New("write fail")}
-	_, err := runStreamingChat(context.Background(), client, nil, nil, w)
+	_, err := runChatRequest(context.Background(), client, requestArgs{}, nil, w)
 	if err == nil || !strings.Contains(err.Error(), "write fail") {
 		t.Fatalf("expected write error, got %v", err)
 	}
@@ -298,7 +298,7 @@ func TestEffectiveModel_Whitespace(t *testing.T) {
 func TestRunSimpleChat_WriteError(t *testing.T) {
 	client := &fakeClient{name: "p", model: "m", resp: "ok"}
 	w := errWriter{err: errors.New("write fail")}
-	_, err := runSimpleChat(context.Background(), client, nil, nil, w)
+	_, err := runChatRequest(context.Background(), client, requestArgs{}, nil, w)
 	if err == nil || !strings.Contains(err.Error(), "write fail") {
 		t.Fatalf("expected write error, got %v", err)
 	}
@@ -448,7 +448,7 @@ func TestRunSimpleChat_ChatError(t *testing.T) {
 		chatErr:    fmt.Errorf("chat broken"),
 	}
 	var out bytes.Buffer
-	_, err := runSimpleChat(context.Background(), client, nil, nil, &out)
+	_, err := runChatRequest(context.Background(), client, requestArgs{}, nil, &out)
 	if err == nil || !strings.Contains(err.Error(), "chat broken") {
 		t.Fatalf("expected chat error, got %v", err)
 	}
diff --git a/internal/lsp/handlers_completion.go b/internal/lsp/handlers_completion.go
index fced629..5ae6cd4 100644
--- a/internal/lsp/handlers_completion.go
+++ b/internal/lsp/handlers_completion.go
@@ -12,6 +12,7 @@ import (
 	"sync"
 	"time"
 
+	"codeberg.org/snonux/hexai/internal/chatrun"
 	"codeberg.org/snonux/hexai/internal/llm"
 	"codeberg.org/snonux/hexai/internal/llmutils"
 	"codeberg.org/snonux/hexai/internal/logging"
@@ -326,12 +327,11 @@ func (cs *completionService) runCompletionForSpec(ctx context.Context, plan comp
 func (cs *completionService) executeChatCompletion(ctx context.Context, plan completionPlan, spec requestSpec, client llm.Client, sortPrefix string) ([]CompletionItem, bool) {
 	s := cs.srv
 	messages := cs.buildCompletionMessages(plan.inlinePrompt, plan.hasExtra, plan.extraText, plan.inParams, plan.params, plan.above, plan.current, plan.below, plan.funcCtx)
-	sentSize := 0
-	for _, m := range messages {
-		sentSize += len(m.Content)
-	}
+	sentSize := chatrun.SentBytes(messages)
 	s.incSentCounters(sentSize)
-	text, err := client.Chat(ctx, messages, spec.options...)
+	// Completion never streams to a writer; Invoke collects the full text and
+	// keeps this path aligned with chatWithStats and the other surfaces.
+	text, err := chatrun.Invoke(ctx, client, messages, spec.options, nil)
 	if err != nil {
 		logging.Logf("lsp ", "llm completion error: %v", err)
 		return nil, false
diff --git a/internal/lsp/handlers_utils.go b/internal/lsp/handlers_utils.go
index a8e87e4..55affaf 100644
--- a/internal/lsp/handlers_utils.go
+++ b/internal/lsp/handlers_utils.go
@@ -10,6 +10,7 @@ import (
 	"unicode/utf8"
 
 	"codeberg.org/snonux/hexai/internal/appconfig"
+	"codeberg.org/snonux/hexai/internal/chatrun"
 	"codeberg.org/snonux/hexai/internal/llm"
 	"codeberg.org/snonux/hexai/internal/llmutils"
 	"codeberg.org/snonux/hexai/internal/logging"
@@ -246,39 +247,34 @@ func isIdentChar(ch byte) bool {
 	return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_'
 }
 
-// chatWithStats wraps llmClient.Chat to increment counters and emit a tmux heartbeat.
+// chatWithStats wraps the LLM request to increment counters and emit a tmux
+// heartbeat. The byte counting and the actual LLM invocation are delegated to
+// the shared chatrun package so the LSP stays in lock-step with the CLI and
+// code-action surfaces; the debounce/throttle gating, in-memory counters and
+// stats-error logging remain here because they are LSP-specific.
 func (s *Server) chatWithStats(ctx context.Context, surface surfaceKind, spec requestSpec, msgs []llm.Message) (string, error) {
-	// Count bytes sent
-	sent := 0
-	for _, m := range msgs {
-		sent += len(m.Content)
-	}
+	sent := chatrun.SentBytes(msgs)
 	s.incSentCounters(sent)
 	// Debounce/throttle if configured (reuse completion gates)
 	s.completionSvc().waitForDebounce(ctx)
 	if !s.waitForThrottle(ctx) {
 		return "", context.Canceled
 	}
-	// Perform request
+	// Resolve the client for this surface/spec.
 	client := s.clientFor(spec)
 	if client == nil {
-		provider := strings.TrimSpace(spec.provider)
-		if provider == "" {
-			provider = strings.TrimSpace(s.currentConfig().Provider)
-		}
-		if provider == "" {
-			return "", fmt.Errorf("llm client unavailable; check the configured provider and required API key")
-		}
-		return "", fmt.Errorf("llm client unavailable for provider %q; check the configured provider and required API key", provider)
+		return "", s.unavailableClientError(spec)
 	}
 	modelUsed := spec.effectiveModel(client.DefaultModel())
-	txt, err := client.Chat(ctx, msgs, spec.options...)
+	// chatWithStats never streams to a writer; pass nil out. Invoke prefers
+	// streaming providers but collects the full text either way.
+	txt, err := chatrun.Invoke(ctx, client, msgs, spec.options, nil)
 	if err != nil {
 		s.logLLMStats(modelUsed)
 		return "", err
 	}
 	s.incRecvCounters(len(txt))
-	// Update global stats cache; log but don't fail on stats errors
+	// Update global stats cache; log but don't fail on stats errors.
 	if err := stats.Update(ctx, client.Name(), modelUsed, sent, len(txt)); err != nil {
 		logging.Logf("lsp ", "stats update error: %v", err)
 	}
@@ -286,6 +282,19 @@ func (s *Server) chatWithStats(ctx context.Context, surface surfaceKind, spec re
 	return txt, nil
 }
 
+// unavailableClientError builds the descriptive error returned when no LLM
+// client could be resolved for spec, naming the configured provider when known.
+func (s *Server) unavailableClientError(spec requestSpec) error {
+	provider := strings.TrimSpace(spec.provider)
+	if provider == "" {
+		provider = strings.TrimSpace(s.currentConfig().Provider)
+	}
+	if provider == "" {
+		return fmt.Errorf("llm client unavailable; check the configured provider and required API key")
+	}
+	return fmt.Errorf("llm client unavailable for provider %q; check the configured provider and required API key", provider)
+}
+
 // Inline prompt utilities
 
 func lineHasInlinePrompt(line string, openStr string, open, close byte) bool {