Improve multi-provider completion streaming and CLI selector flags

8 files changed, 118 insertions, 38 deletions

diff --git a/internal/lsp/chat_trigger_suppression_test.go b/internal/lsp/chat_trigger_suppression_test.go
index 9f9f5bc..852f955 100644
--- a/internal/lsp/chat_trigger_suppression_test.go
+++ b/internal/lsp/chat_trigger_suppression_test.go
@@ -13,7 +13,7 @@ func TestCompletionSuppressedOnChatTriggerEOL(t *testing.T) {
 	tests := []string{"What now?>", "Explain!>", "Refactor:>", "note ;>"}
 	for i, line := range tests {
 		p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://chat-suppr.go"}}
-		items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+		items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 		if !ok {
 			t.Fatalf("case %d: expected ok=true", i)
 		}
diff --git a/internal/lsp/completion_cache_test.go b/internal/lsp/completion_cache_test.go
index 057b5c5..ff85906 100644
--- a/internal/lsp/completion_cache_test.go
+++ b/internal/lsp/completion_cache_test.go
@@ -25,7 +25,7 @@ func TestCompletionCache_IgnoresWhitespaceBeforeCursor(t *testing.T) {
 	// First request with trailing spaces before cursor
 	line := "foo   "
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://x.go"}}
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok || len(items) == 0 || fake.calls != 1 {
 		t.Fatalf("expected first call to invoke LLM; ok=%v len=%d calls=%d", ok, len(items), fake.calls)
 	}
@@ -33,7 +33,7 @@ func TestCompletionCache_IgnoresWhitespaceBeforeCursor(t *testing.T) {
 	// Same logical context but with a different amount of trailing whitespace
 	line2 := "foo             "
 	p2 := CompletionParams{Position: Position{Line: 0, Character: len(line2)}, TextDocument: TextDocumentIdentifier{URI: "file://x.go"}}
-	items2, ok2 := s.tryLLMCompletion(p2, "", line2, "", "", "", false, "")
+	items2, ok2, _ := s.tryLLMCompletion(p2, "", line2, "", "", "", false, "")
 	if !ok2 || len(items2) == 0 {
 		t.Fatalf("expected cache hit to still return items")
 	}
diff --git a/internal/lsp/completion_codex_path_test.go b/internal/lsp/completion_codex_path_test.go
index ea27c6e..6ee8c97 100644
--- a/internal/lsp/completion_codex_path_test.go
+++ b/internal/lsp/completion_codex_path_test.go
@@ -48,7 +48,7 @@ func TestTryLLMCompletion_PrefersCodeCompleterOverChat(t *testing.T) {
 	s.llmClient = fake
 	line := "obj."
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://x.go"}}
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok || len(items) == 0 {
 		t.Fatalf("expected completion items via CodeCompleter path")
 	}
@@ -70,7 +70,7 @@ func TestTryLLMCompletion_FallsBackToChatOnCodeCompleterError(t *testing.T) {
 	s.llmClient = fake
 	line := "obj."
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://y.go"}}
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true even on fallback path")
 	}
diff --git a/internal/lsp/completion_prefix_strip_test.go b/internal/lsp/completion_prefix_strip_test.go
index 6173d6f..e0c655c 100644
--- a/internal/lsp/completion_prefix_strip_test.go
+++ b/internal/lsp/completion_prefix_strip_test.go
@@ -52,7 +52,7 @@ func TestTryLLMCompletion_ManualInvokeAfterWhitespace_Allows(t *testing.T) {
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://x.go"}}
 	// Simulate manual user invocation (TriggerKind=1)
 	p.Context = json.RawMessage([]byte(`{"triggerKind":1}`))
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true for manual invoke after whitespace")
 	}
@@ -72,7 +72,7 @@ func TestTryLLMCompletion_InlinePromptAlwaysTriggers(t *testing.T) {
 	line := "prefix >do something> suffix"
 	// No trigger char immediately before cursor; place cursor at end
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://inline.go"}}
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok || len(items) == 0 {
 		t.Fatalf("expected completion to trigger on inline >text> prompt")
 	}
@@ -89,7 +89,7 @@ func TestTryLLMCompletion_DoubleOpenEmpty_DoesNotAutoTrigger(t *testing.T) {
 	s.llmClient = fake
 	line := ">>   " // empty content after double-open should not force-trigger
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://empty-inline.go"}}
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true for non-trigger path")
 	}
@@ -128,7 +128,7 @@ func TestBareDoubleOpenPreventsAutoTriggerEvenWithOtherTriggers(t *testing.T) {
 	// Place a '.' earlier but also include bare double-open at end; should not auto-trigger
 	line := "obj. call >>"
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://bare-ds.go"}}
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true (handled), but not auto-triggering")
 	}
@@ -152,7 +152,7 @@ func TestBareDoubleOpenOnNextLine_PreventsAutoTrigger(t *testing.T) {
 	current := "expression := flag.String(\"expression\", \"\", \"Expression to evaluate\")"
 	below := ">>"
 	p := CompletionParams{Position: Position{Line: 0, Character: len(current)}, TextDocument: TextDocumentIdentifier{URI: "file://nextline.go"}}
-	items, ok := s.tryLLMCompletion(p, "", current, below, "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", current, below, "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true handled")
 	}
@@ -177,7 +177,7 @@ func TestBareDoubleOpenPreventsManualInvoke(t *testing.T) {
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://bare-ds-manual.go"}}
 	// Simulate manual invoke
 	p.Context = json.RawMessage([]byte(`{"triggerKind":1}`))
-	items, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	items, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true (handled)")
 	}
diff --git a/internal/lsp/debounce_throttle_test.go b/internal/lsp/debounce_throttle_test.go
index 81a2c1a..7efd439 100644
--- a/internal/lsp/debounce_throttle_test.go
+++ b/internal/lsp/debounce_throttle_test.go
@@ -37,7 +37,7 @@ func TestCompletionDebounce_WaitsUntilQuiet(t *testing.T) {
 	p.Context = json.RawMessage([]byte(`{"triggerKind":1}`))
 
 	start := time.Now()
-	_, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
+	_, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, "")
 	if !ok {
 		t.Fatalf("expected ok=true")
 	}
@@ -65,7 +65,7 @@ func TestCompletionThrottle_SerializesCalls(t *testing.T) {
 	p := CompletionParams{Position: Position{Line: 0, Character: len(line)}, TextDocument: TextDocumentIdentifier{URI: "file://throttle.go"}}
 	p.Context = json.RawMessage([]byte(`{"triggerKind":1}`))
 	start := time.Now()
-	if _, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, ""); !ok {
+	if _, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, ""); !ok {
 		t.Fatalf("first call expected ok=true")
 	}
 	if f1.t.IsZero() {
@@ -77,7 +77,7 @@ func TestCompletionThrottle_SerializesCalls(t *testing.T) {
 	s.compCache = make(map[string]string)
 	f2 := &timeLLM{}
 	s.llmClient = f2
-	if _, ok := s.tryLLMCompletion(p, "", line, "", "", "", false, ""); !ok {
+	if _, ok, _ := s.tryLLMCompletion(p, "", line, "", "", "", false, ""); !ok {
 		t.Fatalf("second call expected ok=true")
 	}
 	if f2.t.IsZero() {
diff --git a/internal/lsp/handlers_completion.go b/internal/lsp/handlers_completion.go
index 237d34d..78e685a 100644
--- a/internal/lsp/handlers_completion.go
+++ b/internal/lsp/handlers_completion.go
@@ -45,9 +45,9 @@ func (s *Server) handleCompletion(req Request) {
 		if s.llmClient != nil {
 			newFunc := s.isDefiningNewFunction(p.TextDocument.URI, p.Position)
 			extra, has := s.buildAdditionalContext(newFunc, p.TextDocument.URI, p.Position)
-			items, ok := s.tryLLMCompletion(p, above, current, below, funcCtx, docStr, has, extra)
+			items, ok, incomplete := s.tryLLMCompletion(p, above, current, below, funcCtx, docStr, has, extra)
 			if ok {
-				s.reply(req.ID, CompletionList{IsIncomplete: false, Items: items}, nil)
+				s.reply(req.ID, CompletionList{IsIncomplete: incomplete, Items: items}, nil)
 				return
 			}
 		}
@@ -87,28 +87,33 @@ func (s *Server) logCompletionContext(p CompletionParams, above, current, below,
 		p.TextDocument.URI, p.Position.Line, p.Position.Character, trimLen(above), trimLen(current), trimLen(below), trimLen(funcCtx))
 }
 
-func (s *Server) tryLLMCompletion(p CompletionParams, above, current, below, funcCtx, docStr string, hasExtra bool, extraText string) ([]CompletionItem, bool) {
+func (s *Server) tryLLMCompletion(p CompletionParams, above, current, below, funcCtx, docStr string, hasExtra bool, extraText string) ([]CompletionItem, bool, bool) {
 	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
-	defer cancel()
+	var cancelOnce sync.Once
+	end := func() { cancelOnce.Do(cancel) }
 
 	plan, items, handled := s.prepareCompletionPlan(p, above, current, below, funcCtx, docStr, hasExtra, extraText)
 	if handled {
-		return items, true
+		end()
+		return items, true, false
 	}
 	specs := s.buildRequestSpecs(surfaceCompletion)
 	if len(specs) == 0 {
-		return nil, false
+		end()
+		return nil, false, false
 	}
 	type jobResult struct {
 		items []CompletionItem
 		ok    bool
 	}
-	results := make([]jobResult, len(specs))
+	results := make(chan jobResult, len(specs))
 	var wg sync.WaitGroup
-	var mu sync.Mutex
+	started := 0
 	s.waitForDebounce(ctx)
 	if !s.waitForThrottle(ctx) {
-		return nil, false
+		end()
+		close(results)
+		return nil, false, false
 	}
 	for _, spec := range specs {
 		spec := spec
@@ -116,27 +121,67 @@ func (s *Server) tryLLMCompletion(p CompletionParams, above, current, below, fun
 		if client == nil {
 			continue
 		}
+		started++
 		wg.Add(1)
 		go func(idx int, spec requestSpec, client llm.Client) {
 			defer wg.Done()
 			items, ok := s.runCompletionForSpec(ctx, plan, spec, client)
-			mu.Lock()
-			results[idx] = jobResult{items: items, ok: ok}
-			mu.Unlock()
+			results <- jobResult{items: items, ok: ok}
 		}(spec.index, spec, client)
 	}
-	wg.Wait()
-	accumulated := make([]CompletionItem, 0)
-	for _, res := range results {
-		if !res.ok {
-			continue
+
+	if started == 0 {
+		end()
+		close(results)
+		return nil, false, false
+	}
+
+	go func() {
+		wg.Wait()
+		close(results)
+	}()
+
+	if started == 1 {
+		res := <-results
+		if !res.ok || len(res.items) == 0 {
+			end()
+			return nil, false, false
 		}
-		accumulated = append(accumulated, res.items...)
+		end()
+		return res.items, true, false
 	}
-	if len(accumulated) == 0 {
-		return nil, false
+
+	firstCh := make(chan []CompletionItem, 1)
+	go func(planKey string) {
+		defer end()
+		combined := make([]CompletionItem, 0)
+		firstSent := false
+		for res := range results {
+			if !res.ok || len(res.items) == 0 {
+				continue
+			}
+			combined = append(combined, res.items...)
+			if !firstSent {
+				first := make([]CompletionItem, len(res.items))
+				copy(first, res.items)
+				firstCh <- first
+				firstSent = true
+			}
+		}
+		if !firstSent {
+			close(firstCh)
+			return
+		}
+		s.storePendingCompletion(planKey, combined)
+		close(firstCh)
+	}(plan.cacheKey)
+
+	firstItems, ok := <-firstCh
+	if !ok || len(firstItems) == 0 {
+		end()
+		return nil, false, false
 	}
-	return accumulated, true
+	return firstItems, true, true
 }
 
 func (s *Server) prepareCompletionPlan(p CompletionParams, above, current, below, funcCtx, docStr string, hasExtra bool, extraText string) (completionPlan, []CompletionItem, bool) {
@@ -162,6 +207,9 @@ func (s *Server) prepareCompletionPlan(p CompletionParams, above, current, below
 	plan.inParams = inParamList(current, p.Position.Character)
 	plan.manualInvoke = parseManualInvoke(p.Context)
 	plan.cacheKey = s.completionCacheKey(p, above, current, below, funcCtx, plan.inParams, hasExtra, extraText)
+	if pending := s.takePendingCompletion(plan.cacheKey); len(pending) > 0 {
+		return plan, pending, true
+	}
 	if isBareDoubleOpen(current, openChar, closeChar) || isBareDoubleOpen(below, openChar, closeChar) {
 		logging.Logf("lsp ", "%scompletion skip=empty-double-semicolon line=%d char=%d current=%q%s", logging.AnsiYellow, p.Position.Line, p.Position.Character, trimLen(current), logging.AnsiBase)
 		return plan, []CompletionItem{}, true
diff --git a/internal/lsp/handlers_document.go b/internal/lsp/handlers_document.go
index 9325877..da7db51 100644
--- a/internal/lsp/handlers_document.go
+++ b/internal/lsp/handlers_document.go
@@ -231,7 +231,7 @@ func (s *Server) runInlinePrompt(uri string, pos Position) {
 	docStr := s.buildDocString(p, above, current, below, funcCtx)
 	newFunc := s.isDefiningNewFunction(uri, p.Position)
 	extra, hasExtra := s.buildAdditionalContext(newFunc, uri, p.Position)
-	items, ok := s.tryLLMCompletion(p, above, current, below, funcCtx, docStr, hasExtra, extra)
+	items, ok, _ := s.tryLLMCompletion(p, above, current, below, funcCtx, docStr, hasExtra, extra)
 	if !ok || len(items) == 0 {
 		return
 	}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 1fbb0cc..f8b328b 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -40,8 +40,9 @@ type Server struct {
 	llmRespBytesTotal int64
 	startTime         time.Time
 	// Small LRU cache for recent code completion outputs (keyed by context)
-	compCache      map[string]string
-	compCacheOrder []string // most-recent at end; cap ~10
+	compCache          map[string]string
+	compCacheOrder     []string // most-recent at end; cap ~10
+	pendingCompletions map[string][]CompletionItem
 	// Outgoing JSON-RPC id counter for server-initiated requests
 	nextID      int64
 	lastLLMCall time.Time
@@ -112,6 +113,7 @@ func NewServer(r io.Reader, w io.Writer, logger *log.Logger, opts ServerOptions)
 	s := &Server{in: bufio.NewReader(r), out: w, logger: logger, docs: make(map[string]*document), logContext: opts.LogContext, configStore: opts.ConfigStore}
 	s.startTime = time.Now()
 	s.compCache = make(map[string]string)
+	s.pendingCompletions = make(map[string][]CompletionItem)
 	s.applyOptions(opts)
 	// Initialize dispatch table
 	s.handlers = map[string]func(Request){
@@ -315,6 +317,36 @@ func (s *Server) currentConfig() appconfig.App {
 	return s.cfg
 }
 
+func (s *Server) storePendingCompletion(key string, items []CompletionItem) {
+	if len(items) == 0 {
+		return
+	}
+	cpy := make([]CompletionItem, len(items))
+	copy(cpy, items)
+	s.mu.Lock()
+	if s.pendingCompletions == nil {
+		s.pendingCompletions = make(map[string][]CompletionItem)
+	}
+	s.pendingCompletions[key] = cpy
+	s.mu.Unlock()
+}
+
+func (s *Server) takePendingCompletion(key string) []CompletionItem {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if len(s.pendingCompletions) == 0 {
+		return nil
+	}
+	items, ok := s.pendingCompletions[key]
+	if !ok {
+		return nil
+	}
+	delete(s.pendingCompletions, key)
+	cpy := make([]CompletionItem, len(items))
+	copy(cpy, items)
+	return cpy
+}
+
 func (s *Server) maxTokens() int {
 	cfg := s.currentConfig()
 	if cfg.MaxTokens <= 0 {