diff --git a/internal/loki/chat_conversation.go b/internal/loki/chat_conversation.go index cffb577..a7f64e9 100644 --- a/internal/loki/chat_conversation.go +++ b/internal/loki/chat_conversation.go @@ -840,6 +840,10 @@ func (c *Conversation) generate(ctx context.Context, caps Caps, temperature floa // ne l'ajoute que quand il est là pour ne pas gonfler chaque événement. if ev.ToolUsed.Body != "" { tu["body"] = ev.ToolUsed.Body + tu["body_lines"] = ev.ToolUsed.BodyLines + if ev.ToolUsed.BodyTail { + tu["body_tail"] = true + } } // Lignes +/- d'une écriture (edit / mémoire) : persistées avec le reste // pour que le diff soit encore là après un rafraîchissement. @@ -1289,6 +1293,31 @@ func (c *Conversation) Subscribe(ctx context.Context, from int, emit func(map[st // Bloque jusqu'à ce que ctx (la connexion HTTP) soit annulé — la génération, elle, // continue indépendamment. emit renvoie false si l'écriture échoue (client parti). func (c *Conversation) SubscribeTail(ctx context.Context, from int, tail int, convID string, emit func(map[string]any) bool) { + c.subscribeSink(ctx, from, tail, convID, tailSink{emit: emit}) +} + +// tailSink : sortie d'un abonné. emit écrit et pousse aussitôt ; queue met en +// tampon (en écrivant d'office un tampon plein) et flush pousse ce qui attend. +// Sans queue/flush, tout passe par emit, un événement à la fois. +type tailSink struct { + emit func(map[string]any) bool + queue func(map[string]any) bool + flush func() bool +} + +// subscribeSink : SubscribeTail, en regroupant le direct. Au réveil, plusieurs +// événements peuvent attendre (un jeton chacun) : une écriture et un flush par +// événement multipliaient les appels système et les trames d'un client en +// retard, ou d'un relais chiffré. Ils partent ensemble, dans l'ordre. +func (c *Conversation) subscribeSink(ctx context.Context, from int, tail int, convID string, sink tailSink) { + emit := sink.emit + queue, flush := sink.queue, sink.flush + if queue == nil { + queue = emit + } + if flush == nil { + flush = func() bool { return true } + } // Réveille les attentes de cond quand la connexion se ferme. go func() { <-ctx.Done() @@ -1448,10 +1477,13 @@ func (c *Conversation) SubscribeTail(ctx context.Context, from int, tail int, co // `lastEmitted` (avant mise à jour) sert de repère : si la compaction de // fin de tour vient de fusionner un bloc que le client suivait en direct, // l'événement fusionné arrive avec un Seq supérieur au sien → `replace`. - if !emit(decorateEvent(ev, lastEmitted)) { + if !queue(decorateEvent(ev, lastEmitted)) { return } } + if !flush() { + return + } c.mu.Lock() } } diff --git a/internal/loki/code_retry.go b/internal/loki/code_retry.go index c62c5a8..6507bc6 100644 --- a/internal/loki/code_retry.go +++ b/internal/loki/code_retry.go @@ -43,22 +43,27 @@ var defaultRetryPatterns = []string{ var ( retryOnce sync.Once retryRes []*regexp.Regexp + // retryCustom : la surcharge `retry_patterns` est active. Ses motifs sont + // quelconques (\A, (?s).*? sur des Ko…) : le balayage fenêtré pendant le + // flux (textualToolCallFrom) ne vaut que pour les motifs embarqués. + retryCustom bool ) func retryPatterns() []*regexp.Regexp { - retryOnce.Do(func() { retryRes = compileRetryPatterns() }) + retryOnce.Do(func() { retryRes, retryCustom = compileRetryPatterns() }) return retryRes } -func compileRetryPatterns() []*regexp.Regexp { +func compileRetryPatterns() ([]*regexp.Regexp, bool) { pats := defaultRetryPatterns + custom := false // Surcharge optionnelle : clé d'état `retry_patterns` = tableau JSON de // regex. Pas d'UI dédiée — c'est un réglage d'expert, posé via // `loki config` ou l'API ; l'embarqué couvre les cas connus. if raw := getStr(bkState, "retry_patterns"); raw != "" { - var custom []string - if json.Unmarshal([]byte(raw), &custom) == nil && len(custom) > 0 { - pats = custom + var over []string + if json.Unmarshal([]byte(raw), &over) == nil && len(over) > 0 { + pats, custom = over, true } } var out []*regexp.Regexp @@ -67,7 +72,7 @@ func compileRetryPatterns() []*regexp.Regexp { out = append(out, re) } } - return out + return out, custom } // textualToolCall dit si le texte final du tour contient un appel d'outil @@ -76,6 +81,79 @@ func textualToolCall(content string) bool { return textualToolCallSnippet(content) != "" } +// retryScanMargin / retryScanMaxBack : marge relue avant la fin du balayage +// précédent, et recul au-delà duquel on retombe sur le balayage complet. +const ( + retryScanMargin = 512 + retryScanMaxBack = 16 << 10 +) + +// textualToolCallFrom : textualToolCall pendant le flux, quand content[:scanned] +// a DÉJÀ été balayé sans rien trouver. Relire toute la réponse à chaque morceau +// coûtait jusqu'à 360 µs par jeton sur 80 Ko — du CPU volé au décodage quand +// des experts tournent sur le processeur. Une correspondance neuve finit +// forcément après scanned : on ne relit que la fin, à partir d'un vrai début de +// ligne (sinon `^` verrait un début de ligne au milieu d'une phrase et +// relancerait un tour parfaitement bon). Motifs de la surcharge : balayage +// complet, rien ne borne leur portée. +func textualToolCallFrom(content string, scanned int) bool { + if content == "" { + return false + } + res := retryPatterns() + start := 0 + if !retryCustom { + start = retryScanStart(content, scanned) + } + w := content[start:] + for _, re := range res { + if re.MatchString(w) { + return true + } + } + return false +} + +// retryScanStart : début de la fenêtre à relire. On part de scanned − marge, +// puis on recule tant qu'on est dans une suite que les motifs embarqués peuvent +// traverser sur toute sa longueur (\w+, \s*, `{"name":`, multioctet) : un +// ` len(content) { + scanned = len(content) + } + p := scanned - retryScanMargin + if p <= 0 { + return 0 + } + floor := scanned - retryScanMaxBack + for p > 0 && retryRunByte(content[p-1]) { + p-- + if p < floor { + return 0 + } + } + i := strings.LastIndexByte(content[:p], '\n') + if i < 0 || i+1 < floor { + return 0 + } + return i + 1 +} + +func retryRunByte(c byte) bool { + switch { + case c >= 0x80, c == '_', c >= '0' && c <= '9', c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z': + return true + } + switch c { + case ' ', '\t', '\n', '\r', '\f', '"', ':', '{': + return true + } + return false +} + // textualToolCallSnippet renvoie l'extrait fautif (le motif reconnu et un peu // de ce qui le suit), "" si aucun. Cité dans la consigne corrective : le // modèle voit CE qu'il a mal écrit au lieu d'une remontrance abstraite. diff --git a/internal/loki/code_verify.go b/internal/loki/code_verify.go index 7f26bd3..f73eb42 100644 --- a/internal/loki/code_verify.go +++ b/internal/loki/code_verify.go @@ -219,6 +219,10 @@ func (c *Conversation) forwardStream(ev StreamEvent, epoch int, isolated bool) { } if ev.ToolUsed.Body != "" { tu["body"] = ev.ToolUsed.Body + tu["body_lines"] = ev.ToolUsed.BodyLines + if ev.ToolUsed.BodyTail { + tu["body_tail"] = true + } } if len(ev.ToolUsed.Diff) > 0 { tu["diff"] = ev.ToolUsed.Diff diff --git a/internal/loki/llm_client.go b/internal/loki/llm_client.go index 15b18a0..edfe99d 100644 --- a/internal/loki/llm_client.go +++ b/internal/loki/llm_client.go @@ -551,7 +551,11 @@ type ToolUsedEvent struct { // mem_add/mem_edit), diffusé ligne à ligne pendant que le modèle le tape, // pour que la bulle se remplisse en direct au lieu de rester figée puis de // s'ouvrir d'un coup. Transitoire : seul Diff (état final) est rejoué. - Body string + // Ce n'est que la FIN du corps (bodyTail) : BodyTail dit que le début a été + // omis, BodyLines donne le nombre réel de lignes pour le « +N ». + Body string + BodyTail bool + BodyLines int // Diff : lignes ajoutées/retirées quand l'outil a MODIFIÉ quelque chose // (edit, mem_add, mem_edit). L'UI les affiche en vert (+) et rouge (-). Diff []DiffLine @@ -672,16 +676,11 @@ func writeBodyKey(tool string) string { return "" } -// previewArg pulls the (possibly incomplete) string value of key out of a -// streaming tool-call arguments JSON, so the UI can show the command being -// typed live. Best-effort: it tolerates a truncated tail and basic escapes. -func previewArg(args, key string) string { - v, _ := previewArgDone(args, key) - return v -} - -// previewArgDone : comme previewArg, et dit si la valeur est COMPLÈTE (guillemet -// fermant reçu) — de quoi agir sur un argument avant la fin du flux. +// previewArgDone pulls the (possibly incomplete) string value of key out of a +// streaming tool-call arguments JSON (best-effort: it tolerates a truncated tail +// and basic escapes), et dit si la valeur est COMPLÈTE (guillemet fermant reçu) +// — de quoi agir sur un argument avant la fin du flux. Pour l'affichage en +// direct, argPreview fait la même lecture sans tout relire à chaque morceau. func previewArgDone(args, key string) (string, bool) { i := strings.Index(args, "\""+key+"\"") if i < 0 { @@ -728,6 +727,154 @@ func previewArgDone(args, key string) (string, bool) { return b.String(), closed } +// argPreview : previewArgDone au fil du flux. Pendant l'écriture d'un fichier, +// previewArg relisait et redécodait tout le JSON des arguments à CHAQUE morceau +// — quadratique, et ce CPU-là manque au décodage quand des experts tournent sur +// le processeur. Ici on retient où on en est et on ne décode que les octets +// neufs. Même règle que previewArgDone : première occurrence de "clé", puis +// premier ':', puis premier '"' ; \r ignoré, échappement inconnu (\uXXXX +// compris) recopié tel quel sans la barre. Seule différence, sur un préfixe : +// une barre oblique finale reste en attente de son second octet au lieu d'être +// écrite — elle n'est pas encore un caractère. +type argPreview struct { + key string + pat string // `"key"` + pos int // octets de args déjà consommés + phase int // 0 clé, 1 ':', 2 '"' ouvrant, 3 valeur, 4 valeur fermée + val strings.Builder + lines int // '\n' de la valeur décodée +} + +func newArgPreview(key string) *argPreview { + return &argPreview{key: key, pat: `"` + key + `"`} +} + +// update consomme la suite de args, qui doit prolonger celui des appels +// précédents (les arguments d'un appel ne font que grandir pendant le flux). +func (a *argPreview) update(args string) { + for a.pos < len(args) && a.phase < 4 { + switch a.phase { + case 0: + // La clé peut être coupée entre deux morceaux : on reprend un peu avant. + s := a.pos - (len(a.pat) - 1) + if s < 0 { + s = 0 + } + i := strings.Index(args[s:], a.pat) + if i < 0 { + a.pos = len(args) + return + } + a.pos, a.phase = s+i+len(a.pat), 1 + case 1, 2: + sep := byte(':') + if a.phase == 2 { + sep = '"' + } + j := strings.IndexByte(args[a.pos:], sep) + if j < 0 { + a.pos = len(args) + return + } + a.pos += j + 1 + a.phase++ + case 3: + seg := args[a.pos:] + k := strings.IndexAny(seg, "\\\"") + if k < 0 { + k = len(seg) + } + a.val.WriteString(seg[:k]) + a.lines += strings.Count(seg[:k], "\n") + a.pos += k + if k == len(seg) { + return + } + if seg[k] == '"' { + a.pos++ + a.phase = 4 + return + } + if k+1 >= len(seg) { + return // échappement coupé : on attend son second octet + } + switch e := seg[k+1]; e { + case 'n': + a.val.WriteByte('\n') + a.lines++ + case 't': + a.val.WriteByte('\t') + case 'r': + default: // '"', '\\', '/', et l'inconnu recopié tel quel + a.val.WriteByte(e) + if e == '\n' { + a.lines++ + } + } + a.pos += 2 + } + } +} + +// value : la valeur décodée jusqu'ici (sans copie). +func (a *argPreview) value() string { return a.val.String() } + +// done : guillemet fermant reçu. +func (a *argPreview) done() bool { return a.phase == 4 } + +// bodyTailLines / bodyTailBytes : ce qu'un événement de frappe montre du corps. +const ( + bodyTailLines = 40 + bodyTailBytes = 4096 +) + +// bodyTail : la fin du corps en cours d'écriture, ses 40 dernières lignes et +// 4 Kio au plus, coupée en début de ligne (en début de caractère pour une ligne +// seule plus longue). cut : le début a été omis. La bulle ne fait défiler que +// la fin de toute façon, et le diff final reste complet. +func bodyTail(s string) (tail string, cut bool) { + lo := len(s) - bodyTailBytes + if lo < 0 { + lo = 0 + } + start, n, end := -1, 0, len(s) + for { + i := strings.LastIndexByte(s[:end], '\n') + if i < 0 { + if lo == 0 { + start = 0 + } + break + } + if i+1 < lo { + break + } + n++ + start = i + 1 + if n == bodyTailLines { + break + } + end = i + } + if start < 0 { + start = lo + for start < len(s) && !utf8.RuneStart(s[start]) { + start++ + } + } + return s[start:], start > 0 +} + +// bodyLineCount : lignes du corps comme les compte l'UI (bodyLineCount en JS) — +// un saut de ligne final n'ouvre pas de ligne de plus. +func bodyLineCount(s string) int { + s = strings.TrimSuffix(s, "\n") + if s == "" { + return 0 + } + return strings.Count(s, "\n") + 1 +} + // StatsEvent carries llama.cpp's per-completion timing (final chunk). type StatsEvent struct { PromptTokens int `json:"prompt_tokens,omitempty"` @@ -1281,7 +1428,15 @@ func runChatTools(ctx context.Context, messages []Message, tools []Tool, tempera engineServed() // le slot porte désormais cette requête (llm_slots.go) } toolCalls := map[int]*ToolCall{} + // argBufs : arguments de chaque appel, accumulés sans recopie. `+=` sur la + // chaîne recopiait tout le JSON à chaque morceau — quadratique sur l'écriture + // d'un gros fichier. cur.Function.Arguments reste la valeur de référence : + // on la recale sur le tampon à chaque morceau (String() ne copie rien). + argBufs := map[int]*strings.Builder{} assistantContent := strings.Builder{} + // retryScanned : longueur de assistantContent déjà balayée sans appel + // textuel (voir textualToolCallFrom). + retryScanned := 0 finishReason := "" // Accumulateur de stats : timings (prefill/decode) puis usage (total prompt) // arrivent sur des chunks séparés ; on émet une copie complète à chaque MAJ @@ -1294,6 +1449,9 @@ func runChatTools(ctx context.Context, messages []Message, tools []Tool, tempera usageGen, genChunks := 0, 0 lastPreview := "" // last command preview emitted (to stream the typing) lastBodyLines := -1 // lignes déjà diffusées du corps en cours d'écriture + // Lecture au fil de l'eau de l'argument affiché (commande, chemin…) et du + // corps d'une écriture : seuls les octets neufs sont décodés. + var labelDec, bodyDec *argPreview // sentAnswer : du texte de réponse ou un outil est déjà parti vers l'UI pour // cette complétion (une reprise le doublerait). sentReasoning : seul du // raisonnement est parti, qu'on sait retirer (DropReasoning). shown : texte @@ -1445,7 +1603,14 @@ func runChatTools(ctx context.Context, messages []Message, tools []Tool, tempera if tc.Function.Name != "" { cur.Function.Name = tc.Function.Name } - cur.Function.Arguments += tc.Function.Arguments + ab := argBufs[idx] + if ab == nil { + ab = &strings.Builder{} + ab.WriteString(cur.Function.Arguments) + argBufs[idx] = ab + } + ab.WriteString(tc.Function.Arguments) + cur.Function.Arguments = ab.String() } // Stream the command being typed: extract the partial value and // emit it whenever it grows, so the UI shows it appear live. @@ -1471,25 +1636,40 @@ func runChatTools(ctx context.Context, messages []Message, tools []Tool, tempera case "git_clone": key = "url" } - p := previewArg(cur.Function.Arguments, key) + if labelDec == nil || labelDec.key != key { + labelDec = newArgPreview(key) + } + labelDec.update(cur.Function.Arguments) + p := labelDec.value() // Corps en cours de frappe pour les outils d'écriture : on le diffuse // à la LIGNE, pas au token. Un événement par token republierait tout le // contenu à chaque fois (coût quadratique, et c'est ce flot qui saturait // le rendu mobile) ; à la ligne, le nombre d'événements est celui du - // fichier et l'animation reste fluide. - body := "" + // fichier et l'animation reste fluide. Et chaque événement ne porte + // que la FIN du corps (bodyTail) : republier tout le fichier à chaque + // ligne restait quadratique — 250 Mo de flux pour 120 Ko écrits. + full := "" if bk := writeBodyKey(cur.Function.Name); bk != "" { - body = previewArg(cur.Function.Arguments, bk) + if bodyDec == nil || bodyDec.key != bk { + bodyDec = newArgPreview(bk) + } + bodyDec.update(cur.Function.Arguments) + full = bodyDec.value() } - grew := body != "" && strings.Count(body, "\n") > lastBodyLines + grew := full != "" && bodyDec.lines > lastBodyLines if (p != "" && p != lastPreview) || grew { if p != "" { lastPreview = p } if grew { - lastBodyLines = strings.Count(body, "\n") + lastBodyLines = bodyDec.lines } - if !scb(StreamEvent{ToolUsed: &ToolUsedEvent{Name: cur.Function.Name, Label: lastPreview, Body: body, Typing: true}}) { + tu := &ToolUsedEvent{Name: cur.Function.Name, Label: lastPreview, Typing: true} + if full != "" { + tu.Body, tu.BodyTail = bodyTail(full) + tu.BodyLines = bodyLineCount(full) + } + if !scb(StreamEvent{ToolUsed: tu}) { aborted = true break } @@ -1541,8 +1721,18 @@ func runChatTools(ctx context.Context, messages []Message, tools []Tool, tempera // faux appel — parfois un fichier entier — qui ne sera jamais // exécuté. La relance corrective part juste après (voir fin de // boucle). Testé seulement quand le morceau peut ouvrir un motif. + // Balayage limité à ce qui suit le précédent (retryScanned, avancé + // seulement quand un balayage a VRAIMENT eu lieu : un morceau sans + // caractère déclencheur n'est pas relu, l'écart doit rester couvert). + // La passe complète de fin de tour, elle, reste entière. + hit := false if (callsOn && patternRetries < maxPatternRetries || toolChoiceNone) && - strings.ContainsAny(ch.Delta.Content, "<`{[_.") && textualToolCall(assistantContent.String()) { + strings.ContainsAny(ch.Delta.Content, "<`{[_.") { + s := assistantContent.String() + hit = textualToolCallFrom(s, retryScanned) + retryScanned = len(s) + } + if hit { // Sous tool_choice « none », le moteur ne parse plus les // appels : le balisage arrive en texte. Ce n'est pas une // réponse, on retombe sur la relance sans outils. diff --git a/internal/loki/relay_e2e.go b/internal/loki/relay_e2e.go index 96ae0de..75dabee 100644 --- a/internal/loki/relay_e2e.go +++ b/internal/loki/relay_e2e.go @@ -164,24 +164,17 @@ func handleE2EChat(w http.ResponseWriter, r *http.Request) { flusher, _ := w.(http.Flusher) mu, stop := sseHeartbeat(w, flusher) defer stop() - emit := func(obj map[string]any) bool { + // Un sceau par événement, même quand le direct les écrit par lots. + frame := func(obj map[string]any) ([]byte, bool) { b, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{"delta": obj}}}) nonce := make([]byte, 12) if _, err := rand.Read(nonce); err != nil { - return false + return nil, false } sealedEv := append(nonce, gcm.Seal(nil, nonce, b, nil)...) - mu.Lock() - defer mu.Unlock() - if _, err := w.Write([]byte("data: " + base64.StdEncoding.EncodeToString(sealedEv) + "\n\n")); err != nil { - return false - } - if flusher != nil { - flusher.Flush() - } - return true + return []byte("data: " + base64.StdEncoding.EncodeToString(sealedEv) + "\n\n"), true } - runChatStream(r.Context(), body, emit) + runChatStream(r.Context(), body, &sseStream{w: w, flusher: flusher, mu: mu, frame: frame}) } // handleE2EReq : proxy de CONTRÔLE chiffré de bout en bout. Même enveloppe que le diff --git a/internal/loki/stream_hotpath_test.go b/internal/loki/stream_hotpath_test.go new file mode 100644 index 0000000..900ee89 --- /dev/null +++ b/internal/loki/stream_hotpath_test.go @@ -0,0 +1,480 @@ +package loki + +import ( + "context" + "encoding/json" + "math/rand" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + "unicode/utf8" +) + +// chunked découpe s en morceaux de 1 à max octets (graine fixe : reproductible). +func chunked(s string, r *rand.Rand, max int) []string { + var out []string + for len(s) > 0 { + n := 1 + r.Intn(max) + if n > len(s) { + n = len(s) + } + out = append(out, s[:n]) + s = s[n:] + } + return out +} + +// checkArgPreview rejoue args morceau par morceau et compare à chaque étape +// argPreview au previewArgDone historique — à une barre oblique finale près, +// qu'argPreview garde en attente au lieu de l'écrire. +func checkArgPreview(t *testing.T, args, key string, chunks []string) { + t.Helper() + a := newArgPreview(key) + prefix := "" + for _, c := range chunks { + prefix += c + a.update(prefix) + want, wantDone := previewArgDone(prefix, key) + got := a.value() + if a.phase == 3 && a.pos == len(prefix)-1 && prefix[a.pos] == '\\' { + got += "\\" + } + if got != want || a.done() != wantDone { + t.Fatalf("args=%q préfixe=%q : argPreview=(%q,%v), previewArgDone=(%q,%v)", args, prefix, got, a.done(), want, wantDone) + } + if n := strings.Count(a.value(), "\n"); n != a.lines { + t.Fatalf("lines=%d, la valeur en compte %d", a.lines, n) + } + } +} + +func TestArgPreviewMatchesPreviewArgDone(t *testing.T) { + r := rand.New(rand.NewSource(7)) + alphabet := []string{"a", "é", "\n", "\t", "\r", "\\", "\"", "<", "&", "
", "/", " ", "x", "content", ":", "{"} + var cases []string + for i := 0; i < 300; i++ { + var b strings.Builder + for j := r.Intn(200); j > 0; j-- { + b.WriteString(alphabet[r.Intn(len(alphabet))]) + } + v := b.String() + m, _ := json.Marshal(map[string]string{"file": "a/b.go", "content": v}) + cases = append(cases, string(m)) + // Ordre inverse et clé citée avant la vraie : même règle de première occurrence. + cases = append(cases, `{"content" : "`+strings.ReplaceAll(v, `"`, `\"`)+`","file":"x"}`) + cases = append(cases, `{"note":"le champ \"content\" vient après","content":"`+strings.ReplaceAll(v, `"`, `\"`)+`"}`) + } + cases = append(cases, + `{"content":"été \\u0041 fin"}`, + `{"content":"ligne1\nligne2\\`, + `{"content":`, `{"cont`, `{"content":"`, `{"content":"x"}`, ``, + ) + for _, args := range cases { + for _, max := range []int{1, 3, 17, 4096} { + checkArgPreview(t, args, "content", chunked(args, r, max)) + } + } +} + +func FuzzArgPreview(f *testing.F) { + f.Add(`{"file":"a","content":"x\ny\\\"zA"}`, int64(1)) + f.Add(`{"content":"a\\`, int64(2)) + f.Add(`{"content" : "x", "content":"y"}`, int64(3)) + f.Fuzz(func(t *testing.T, args string, seed int64) { + checkArgPreview(t, args, "content", chunked(args, rand.New(rand.NewSource(seed)), 9)) + }) +} + +func TestBodyTail(t *testing.T) { + lines := func(n int, w int) string { + var b strings.Builder + for i := 0; i < n; i++ { + b.WriteString(strings.Repeat("x", w) + strconv.Itoa(i) + "\n") + } + return b.String() + } + long := strings.Repeat("é", 5000) // une seule ligne de 10 000 octets + for _, tc := range []struct { + name string + in string + cut bool + maxLine int + }{ + {"vide", "", false, 0}, + {"court", "a\nb\nc", false, 3}, + {"40 lignes pile", lines(39, 3) + "fin", false, 40}, + {"beaucoup de lignes", lines(500, 3) + "fin", true, bodyTailLines}, + {"lignes larges", lines(100, 300), true, bodyTailLines}, + {"ligne seule énorme", long, true, 1}, + } { + got, cut := bodyTail(tc.in) + if cut != tc.cut || !strings.HasSuffix(tc.in, got) { + t.Errorf("%s : cut=%v, suffixe=%v", tc.name, cut, strings.HasSuffix(tc.in, got)) + } + if len(got) > bodyTailBytes { + t.Errorf("%s : %d octets", tc.name, len(got)) + } + if n := strings.Count(got, "\n") + 1; got != "" && n > tc.maxLine { + t.Errorf("%s : %d lignes", tc.name, n) + } + if cut { + if start := len(tc.in) - len(got); tc.in[start-1] != '\n' && !strings.Contains(tc.name, "seule") { + t.Errorf("%s : coupé en milieu de ligne", tc.name) + } + if got != "" && !json.Valid([]byte(strconv.Quote(got))) || strings.ContainsRune(got, '�') { + t.Errorf("%s : caractère coupé", tc.name) + } + } + } +} + +func TestBodyLineCount(t *testing.T) { + for in, want := range map[string]int{"": 0, "\n": 0, "a": 1, "a\n": 1, "a\nb": 2, "a\nb\n": 2, "\n\n": 2} { + if got := bodyLineCount(in); got != want { + t.Errorf("bodyLineCount(%q) = %d, attendu %d", in, got, want) + } + } +} + +// Les événements de frappe d'un gros write ne portent qu'une fin bornée du +// corps, avec le vrai nombre de lignes ; les arguments exécutés restent entiers. +func TestWriteTypingBodyBounded(t *testing.T) { + var content strings.Builder + for i := 0; i < 3000; i++ { + content.WriteString("ligne numéro " + strconv.Itoa(i) + " du fichier\n") + } + withWorkspace(t) + args, _ := json.Marshal(map[string]string{"file": "gros.txt", "content": content.String()}) + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + if atomic.AddInt32(&calls, 1) > 1 { + _, _ = w.Write([]byte(sseChunk("fini"))) + _, _ = w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"stop"}]}` + "\n\ndata: [DONE]\n\n")) + return + } + s := string(args) + for len(s) > 0 { + n := 37 + if n > len(s) { + n = len(s) + } + for n < len(s) && !utf8.RuneStart(s[n]) { + n++ + } + piece, _ := json.Marshal(s[:n]) + s = s[n:] + _, _ = w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"w1","type":"function","function":{"name":"write","arguments":` + string(piece) + `}}]}}]}` + "\n\n")) + } + _, _ = w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\n\ndata: [DONE]\n\n")) + })) + t.Cleanup(srv.Close) + u, _ := url.Parse(srv.URL) + if err := SetConfigKey("PORT", u.Port()); err != nil { + t.Fatal(err) + } + var sent, last int + var lastEv *ToolUsedEvent + _, err := runChat(context.Background(), []Message{{Role: "user", Content: "écris"}}, 0.7, Caps{Agent: true}, func(ev StreamEvent) bool { + if tu := ev.ToolUsed; tu != nil && tu.Typing && tu.Body != "" { + if len(tu.Body) > bodyTailBytes { + t.Fatalf("corps de %d octets dans un événement de frappe", len(tu.Body)) + } + if tu.BodyLines < last { + t.Fatalf("body_lines recule : %d après %d", tu.BodyLines, last) + } + last = tu.BodyLines + sent += len(tu.Body) + lastEv = tu + } + return true + }) + if err != nil { + t.Fatal(err) + } + if lastEv == nil || !lastEv.BodyTail || lastEv.BodyLines < 2990 { + t.Fatalf("dernier événement = %+v", lastEv) + } + if !strings.HasSuffix(content.String(), lastEv.Body) { + t.Fatalf("la fin diffusée n'est pas la fin du corps : %q", lastEv.Body) + } + // Linéaire : 3000 lignes × 4 Kio au plus, loin des centaines de Mo d'avant. + if sent > 3000*bodyTailBytes { + t.Fatalf("%d octets de corps diffusés", sent) + } + // Ce qui est exécuté, lui, est entier. + if b, err := os.ReadFile(filepath.Join(agentCwd(), "gros.txt")); err != nil || string(b) != content.String() { + t.Fatalf("fichier écrit : %d octets, err=%v", len(b), err) + } +} + +// replayScan rejoue le balayage du flux : un morceau sans caractère déclencheur +// n'est pas balayé. Renvoie l'indice du morceau qui coupe le flux, -1 sinon. +func replayScan(chunks []string, windowed bool) int { + var b strings.Builder + scanned := 0 + for i, c := range chunks { + b.WriteString(c) + if !strings.ContainsAny(c, "<`{[_.") { + continue + } + s := b.String() + var hit bool + if windowed { + hit = textualToolCallFrom(s, scanned) + scanned = len(s) + } else { + hit = textualToolCall(s) + } + if hit { + return i + } + } + return -1 +} + +func TestTextualToolCallWindowMatchesFullScan(t *testing.T) { + r := rand.New(rand.NewSource(11)) + prose := []string{ + "Voici la suite de l'analyse, ligne par ligne. ", + "Le fichier contient trois fonctions ; aucune n'est exportée.\n", + "```go\nfunc main() { fmt.Println(\"x\") }\n```\n", + "- point important\n", + "le JSON {\"name\": \"loki\"} de config, ", + "sans ponctuation du tout mais avec des mots très longs ", + "\n\n", + "Un
en HTML. ", + "été à la plage ", + } + filler := func(n int) string { + var b strings.Builder + for b.Len() < n { + b.WriteString(prose[r.Intn(len(prose))]) + } + return b.String() + } + bad := []string{ + "{\"name\":\"bash\"}", + "\n{\"name\": \"bash\", \"arguments\": {}}", + "\n \n {\"name\":" + strings.Repeat(" \n", 400) + "\"mem_" + strings.Repeat("a", 700) + "\"", + "\n" + strings.Repeat(" ", 900) + "{\"name\": \"web_open\"", + "", + " functions." + strings.Repeat("c", 600) + strings.Repeat(" ", 300) + "(", + "Je lance default_api:bash ", + "```tool_code\nprint(1)\n```", + "", + "[TOOL_REQUEST]", + } + neutral := []string{ + " milieu de ligne {\"name\": \"bash\"} ", // `^` : seulement en début de ligne + strings.Repeat("x", 800) + " {\"name\": \"bash\"", + " presque", + } + var texts []string + for i := 0; i < 120; i++ { + head := filler(500 + r.Intn(5000)) + tail := filler(r.Intn(1500)) + switch i % 3 { + case 0: + texts = append(texts, head+bad[r.Intn(len(bad))]+tail) + case 1: + texts = append(texts, head+neutral[r.Intn(len(neutral))]+tail) + default: + texts = append(texts, head+tail) + } + } + // Motif à cheval sur un morceau sans caractère déclencheur. + texts = append(texts, filler(3000)+""+" puis du texte sans rien") + hits, misses := 0, 0 + for _, s := range texts { + for _, max := range []int{3, 23, 200} { + ch := chunked(s, r, max) + if s == texts[len(texts)-1] { + ch = append(chunked(s[:len(s)-len(" puis du texte sans rien")], r, max), "", " puis du texte", " sans rien") + } + w, f := replayScan(ch, true), replayScan(ch, false) + if w != f { + t.Fatalf("fenêtré coupe au morceau %d, balayage complet au %d (taille %d, max %d)", w, f, len(s), max) + } + if f >= 0 { + hits++ + } else { + misses++ + } + } + } + if hits == 0 || misses == 0 { + t.Fatalf("corpus déséquilibré : %d coupures, %d sans", hits, misses) + } +} + +// La fenêtre est bien une fenêtre sur un texte ordinaire, et démarre en début +// de ligne. +func TestRetryScanStartWindows(t *testing.T) { + s := strings.Repeat("Une phrase ordinaire, avec sa ponctuation.\n", 500) + start := retryScanStart(s, len(s)) + if start == 0 || len(s)-start > 2*retryScanMargin { + t.Fatalf("fenêtre de %d octets sur %d", len(s)-start, len(s)) + } + if s[start-1] != '\n' { + t.Fatal("la fenêtre ne démarre pas en début de ligne") + } + // Une ligne sans fin plus longue que le recul permis : balayage complet. + if got := retryScanStart(strings.Repeat("a b ", 20000), 80000); got != 0 { + t.Fatalf("start = %d, attendu 0", got) + } +} + +// Surcharge retry_patterns : ses motifs ne sont pas bornés, donc pas de fenêtre. +func TestTextualToolCallFromCustomFullScan(t *testing.T) { + retryPatterns() + saved := retryRes + savedCustom := retryCustom + defer func() { retryRes, retryCustom = saved, savedCustom }() + retryRes, retryCustom = []*regexp.Regexp{regexp.MustCompile(`(?mi)\Aoups`)}, true + s := "oups.\n" + strings.Repeat("Une phrase, puis une autre.\n", 400) + if !textualToolCallFrom(s, len(s)-10) { + t.Fatal("motif ancré en tête manqué") + } + retryCustom = false + if textualToolCallFrom(s, len(s)-10) { + t.Fatal("la fenêtre aurait dû le manquer (le test ne prouve rien sinon)") + } +} + +// sseStream : un lot = une écriture, les événements dans l'ordre, et un lot +// trop gros est découpé. +type countingRW struct { + *httptest.ResponseRecorder + writes int +} + +func (c *countingRW) Write(b []byte) (int, error) { + c.writes++ + return c.ResponseRecorder.Write(b) +} + +func TestSSEStreamBatches(t *testing.T) { + rw := &countingRW{ResponseRecorder: httptest.NewRecorder()} + s := &sseStream{w: rw, mu: &sync.Mutex{}, frame: func(m map[string]any) ([]byte, bool) { + return []byte("data: " + strconv.Itoa(m["seq"].(int)) + "\n\n"), true + }} + for i := 1; i <= 10; i++ { + if !s.queue(map[string]any{"seq": i}) { + t.Fatal("queue") + } + } + if rw.writes != 0 { + t.Fatalf("%d écritures avant flush", rw.writes) + } + s.flush() + if rw.writes != 1 { + t.Fatalf("%d écritures pour un lot", rw.writes) + } + for i := 11; i <= 10+2*sseBatchEvents; i++ { + s.queue(map[string]any{"seq": i}) + } + s.emit(map[string]any{"seq": 11 + 2*sseBatchEvents}) + if rw.writes != 4 { + t.Fatalf("%d écritures, attendu 4", rw.writes) + } + var want strings.Builder + for i := 1; i <= 11+2*sseBatchEvents; i++ { + want.WriteString("data: " + strconv.Itoa(i) + "\n\n") + } + if rw.Body.String() != want.String() { + t.Fatal("ordre des événements perdu") + } +} + +// Le direct passe par queue puis flush, dans l'ordre des seq. +func TestSubscribeSinkLiveBatchOrder(t *testing.T) { + c := newTestConv() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var mu sync.Mutex + var seqs []int + flushed := 0 + caught := make(chan struct{}, 1) + sink := tailSink{ + emit: func(m map[string]any) bool { + if _, ok := m["caught_up"]; ok { + caught <- struct{}{} + } + return true + }, + queue: func(m map[string]any) bool { + mu.Lock() + if s, ok := m["seq"].(int); ok { + seqs = append(seqs, s) + } + mu.Unlock() + return true + }, + flush: func() bool { mu.Lock(); flushed++; mu.Unlock(); return true }, + } + go c.subscribeSink(ctx, 0, -1, "", sink) + <-caught + const n = 500 + for i := 0; i < n; i++ { + c.appendDelta(c.epoch, map[string]any{"content": "t"}) + } + deadline := time.Now().Add(5 * time.Second) + for { + mu.Lock() + got, f := len(seqs), flushed + mu.Unlock() + if got == n && f > 0 { + break + } + if time.Now().After(deadline) { + t.Fatalf("%d/%d événements reçus", got, n) + } + time.Sleep(5 * time.Millisecond) + } + mu.Lock() + defer mu.Unlock() + for i, s := range seqs { + if s != i+1 { + t.Fatalf("seq %d en position %d", s, i) + } + } +} + +func BenchmarkWriteTypingPreview(b *testing.B) { + var content strings.Builder + for content.Len() < 120<<10 { + content.WriteString("une ligne de code assez ordinaire, avec \"guillemets\"\n") + } + args, _ := json.Marshal(map[string]string{"file": "x", "content": content.String()}) + s := string(args) + b.ResetTimer() + for i := 0; i < b.N; i++ { + a := newArgPreview("content") + last := -1 + for end := 8; end <= len(s); end += 8 { + a.update(s[:end]) + if a.lines > last { + last = a.lines + bodyTail(a.value()) + } + } + } +} + +func BenchmarkTextualToolCallWindowed(b *testing.B) { + s := strings.Repeat("Une phrase ordinaire, avec sa ponctuation. Et un point.\n", 1500) + b.ResetTimer() + for i := 0; i < b.N; i++ { + textualToolCallFrom(s, len(s)-8) + } +} diff --git a/internal/loki/ui/index.html b/internal/loki/ui/index.html index d31611b..deb7adf 100644 --- a/internal/loki/ui/index.html +++ b/internal/loki/ui/index.html @@ -6900,7 +6900,9 @@ function renderToolMsg(el, tu){ let add=0, del=0; if(tu.added!=null || tu.removed!=null){ add=tu.added||0; del=tu.removed||0; } else if(tu.diff && tu.diff.length){ tu.diff.forEach(l=>{ if(l.op==='+') add++; else if(l.op==='-') del++; }); } - else if(tu.body){ add=bodyLineCount(tu.body); } + // Pendant la frappe, body n'est que la FIN du corps : body_lines porte le + // vrai total (repli sur le décompte pour un serveur d'avant ce champ). + else if(tu.body){ add=tu.body_lines!=null ? tu.body_lines : bodyLineCount(tu.body); } if(add||del) setLabelCounts(el, add, del); const body=bodyOf(el); body.innerHTML=''; const head=document.createElement('div'); head.className='tool-head'; @@ -6923,6 +6925,11 @@ function renderToolMsg(el, tu){ sub.textContent='écriture en cours'; // le +N vit sur l'étiquette (visible repliée) body.appendChild(sub); const pre=document.createElement('pre'); pre.className='diff live'; + // Début omis par le serveur (seules les dernières lignes voyagent) : on le dit. + if(tu.body_tail){ + const more=document.createElement('span'); more.className='dl'; more.textContent=' …'; + pre.appendChild(more); + } lines.forEach((t,i)=>{ const ln=document.createElement('span'); ln.className='dl add'+(i===lines.length-1?' fresh':''); diff --git a/internal/loki/ui/src/js/08-chat-render.js b/internal/loki/ui/src/js/08-chat-render.js index 2d0c908..1275ede 100644 --- a/internal/loki/ui/src/js/08-chat-render.js +++ b/internal/loki/ui/src/js/08-chat-render.js @@ -317,7 +317,9 @@ function renderToolMsg(el, tu){ let add=0, del=0; if(tu.added!=null || tu.removed!=null){ add=tu.added||0; del=tu.removed||0; } else if(tu.diff && tu.diff.length){ tu.diff.forEach(l=>{ if(l.op==='+') add++; else if(l.op==='-') del++; }); } - else if(tu.body){ add=bodyLineCount(tu.body); } + // Pendant la frappe, body n'est que la FIN du corps : body_lines porte le + // vrai total (repli sur le décompte pour un serveur d'avant ce champ). + else if(tu.body){ add=tu.body_lines!=null ? tu.body_lines : bodyLineCount(tu.body); } if(add||del) setLabelCounts(el, add, del); const body=bodyOf(el); body.innerHTML=''; const head=document.createElement('div'); head.className='tool-head'; @@ -340,6 +342,11 @@ function renderToolMsg(el, tu){ sub.textContent='écriture en cours'; // le +N vit sur l'étiquette (visible repliée) body.appendChild(sub); const pre=document.createElement('pre'); pre.className='diff live'; + // Début omis par le serveur (seules les dernières lignes voyagent) : on le dit. + if(tu.body_tail){ + const more=document.createElement('span'); more.className='dl'; more.textContent=' …'; + pre.appendChild(more); + } lines.forEach((t,i)=>{ const ln=document.createElement('span'); ln.className='dl add'+(i===lines.length-1?' fresh':''); diff --git a/internal/loki/web_chat.go b/internal/loki/web_chat.go index e33f4aa..7a1f5f0 100644 --- a/internal/loki/web_chat.go +++ b/internal/loki/web_chat.go @@ -83,12 +83,76 @@ func sseHeartbeat(w http.ResponseWriter, flusher http.Flusher) (*sync.Mutex, fun // se ferme. La GÉNÉRATION est lancée séparément par /api/chat/send dans une // goroutine détachée — fermer le navigateur n'arrête donc plus rien. Partagé par // handleChat (clair) et handleE2EChat (chiffré). -func runChatStream(ctx context.Context, body chatReq, emit func(map[string]any) bool) { +func runChatStream(ctx context.Context, body chatReq, out *sseStream) { tail := -1 if body.Tail != nil { tail = *body.Tail } - conv.SubscribeTail(ctx, body.From, tail, body.ConvID, emit) + conv.subscribeSink(ctx, body.From, tail, body.ConvID, tailSink{emit: out.emit, queue: out.queue, flush: out.flush}) +} + +// sseBatchBytes / sseBatchEvents : taille d'un envoi groupé. Un client très en +// retard peut trouver des milliers d'événements en attente : on écrit par +// tranches au lieu de tout empiler en mémoire. +const ( + sseBatchBytes = 128 << 10 + sseBatchEvents = 256 +) + +// sseStream : écriture des trames `data:` d'un flux d'abonnement. frame met un +// événement en forme (en clair, ou scellé pour le relais E2E — un sceau par +// événement, le client déchiffre chaque ligne `data:` séparément). mu est +// celui du battement de cœur : jamais deux écritures mêlées sur w. +type sseStream struct { + w http.ResponseWriter + flusher http.Flusher + mu *sync.Mutex + frame func(map[string]any) ([]byte, bool) + buf []byte + n int +} + +// emit : un événement, écrit et poussé tout de suite (après ce qui attend). +func (s *sseStream) emit(obj map[string]any) bool { + return s.queue(obj) && s.flush() +} + +// queue : un événement de plus dans le tampon, écrit d'office s'il déborde. +func (s *sseStream) queue(obj map[string]any) bool { + b, ok := s.frame(obj) + if !ok { + return false + } + s.buf = append(s.buf, b...) + s.n++ + if len(s.buf) >= sseBatchBytes || s.n >= sseBatchEvents { + return s.flush() + } + return true +} + +// flush : une seule écriture et un seul flush pour tout ce qui attend. +func (s *sseStream) flush() bool { + if len(s.buf) == 0 { + return true + } + s.mu.Lock() + defer s.mu.Unlock() + _, err := s.w.Write(s.buf) + // Un gros événement (diff final) ne doit pas garder son tampon à vie. + if cap(s.buf) > 4*sseBatchBytes { + s.buf = nil + } else { + s.buf = s.buf[:0] + } + s.n = 0 + if err != nil { + return false + } + if s.flusher != nil { + s.flusher.Flush() + } + return true } // handleChatSend ajoute un message et lance la génération en arrière-plan. Réponse @@ -244,17 +308,9 @@ func handleChat(w http.ResponseWriter, r *http.Request) { flusher, _ := w.(http.Flusher) mu, stop := sseHeartbeat(w, flusher) defer stop() - emit := func(obj map[string]any) bool { + frame := func(obj map[string]any) ([]byte, bool) { b, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{"delta": obj}}}) - mu.Lock() - defer mu.Unlock() - if _, err := w.Write([]byte("data: " + string(b) + "\n\n")); err != nil { - return false - } - if flusher != nil { - flusher.Flush() - } - return true + return []byte("data: " + string(b) + "\n\n"), true } - runChatStream(r.Context(), body, emit) + runChatStream(r.Context(), body, &sseStream{w: w, flusher: flusher, mu: mu, frame: frame}) }