diff --git a/README.md b/README.md index 7529515..0c33076 100644 --- a/README.md +++ b/README.md @@ -227,7 +227,7 @@ redémarrages (volume `/data`). Variables d'environnement du conteneur : | `LOKI_AUTO_UPDATE` | `0` : pas de mise à jour des binaires au démarrage | `1` | | `LOKI_FACADE` | `0` : sans façade Rust, le Go sert directement le port public (porte de sortie) | `1` | | `LOKI_GO_PORT` | port local (127.0.0.1) du serveur Go derrière la façade | `8091` | -| `LOKI_CHAT` | qui tient la conversation : la façade Rust par défaut, outils fichiers, shell et git compris, avec l'export, le mode Code et les critères de la discussion ; le Go garde les autres outils et le reste du mode Code. La façade sonde aussi le gabarit du moteur après ses complétions (sonde de gabarit, `NUDGE_IN_TOOL`), range la télémétrie de ses complétions et ses requêtes en vol auprès du Go, honore les motifs de relance personnalisés (`retry_patterns`) et les slots du moteur (`SIDE_SLOT`, `SLOT_PERSIST`, isolation après la vérification) et les images gardées (`KEEP_TURN_IMAGES`). `go` : la rend au Go (porte de sortie). Avec `PROJ_SNAPSHOT` ou `PREWARM`, le Go la garde de toute façon (décidé au démarrage) | Rust | +| `LOKI_CHAT` | qui tient la conversation : la façade Rust par défaut, outils fichiers, shell et git compris, avec l'export, le mode Code et les critères de la discussion ; le Go garde les autres outils et le reste du mode Code. La façade sonde aussi le gabarit du moteur après ses complétions (sonde de gabarit, `NUDGE_IN_TOOL`), range la télémétrie de ses complétions et ses requêtes en vol auprès du Go, honore les motifs de relance personnalisés (`retry_patterns`) et les slots du moteur (`SIDE_SLOT`, `SLOT_PERSIST`, isolation après la vérification) les images gardées (`KEEP_TURN_IMAGES`) et le préchauffage (`PREWARM`). `go` : la rend au Go (porte de sortie). Avec `PROJ_SNAPSHOT`, le Go la garde de toute façon (décidé au démarrage) | Rust | | `LOKI_MODELS` | `rust` : les routes des modèles — liste, dossiers, suppression, téléchargements, catalogue Hugging Face — sont servies par la façade Rust (expérimental, migration en cours) | Go | | `LOKI_DICTATE` | `rust` : la dictée vocale (`/api/transcribe`, `/api/dictate/*`) et son serveur de transcription sont tenus par la façade Rust (expérimental, migration en cours) | Go | | `LOKI_SETTINGS` | `rust` : les réglages partagés — préférences d'affichage, prompt système, intensité du raisonnement, adresse publique — sont servis par la façade Rust (expérimental, migration en cours) | Go | diff --git a/docs/superpowers/specs/2026-10-16-rust-migration-lot33-prechauffage-design.md b/docs/superpowers/specs/2026-10-16-rust-migration-lot33-prechauffage-design.md new file mode 100644 index 0000000..158cded --- /dev/null +++ b/docs/superpowers/specs/2026-10-16-rust-migration-lot33-prechauffage-design.md @@ -0,0 +1,112 @@ +# Lot 33 — Préchauffage (PREWARM) tenu par la façade + +## Contexte + +Certains recalculs sont inévitables : + +- le prompt réécrit par une compaction ; +- le dernier message assistant rendu autrement au tour suivant ; +- la discussion reprise après une tâche qui a pris le slot. + +Sans la clé, ils tombent sur le premier jeton du message suivant. + +`PREWARM` (`chat_prewarm.go`) envoie au moteur local, dès qu'il est libre, la +requête du prochain tour telle qu'elle partira, suivie d'une sentinelle « . » +et avec `max_tokens` 1. Le moteur calcule tout le préfixe ; le vrai tour ne +calcule plus que son propre message. + +Le préchauffage ne passe jamais devant un vrai travail : + +- toute requête de Loki l'annule, sauf un tour qui le prolonge exactement ; +- un seul est en vol à la fois ; +- aucune erreur ne remonte. + +La clé gardait la conversation au Go. + +## Le lot + +La façade tient l'historique : c'est elle qui assemble et envoie. Le Go +garde le compteur des requêtes en vol (`engineGate`) et ce qui annule. + +### Go (`facade_prewarm.go`) + +- `/internal/v1/prewarm/check` dit si le préchauffage peut partir : + - clé posée, moteur local ; + - ni travail annexe, ni bench, ni optimiseur en cours ; + - loin de minuit ; + - moteur en bonne santé ; + - slot visé : un seul slot, ou celui de la discussion avec SIDE_SLOT. +- `/internal/v1/prewarm/begin` : + - recharge l'état SLOT_PERSIST de la discussion ; + - inscrit le préchauffage comme LA requête en vol (`prewarmBegin`) ; + - son annulation (une autre requête part) est relayée à la façade + (`/internal/facade/v1/chat/prewarm/cancel`). +- `/internal/v1/prewarm/end` : + - rend la place ; + - note ce que le moteur a repris (`noteEnginePrompt`). +- `/internal/v1/oai/begin` accepte `prewarm_keep` : l'étape d'un tour de la + façade qui prolonge le préchauffage ne l'annule pas. +- Après une tâche planifiée, le déclencheur part à la façade + (`/internal/facade/v1/chat/prewarm`). +- PREWARM ne garde plus la conversation. + +### Rust (`chat/prewarm.rs`) + +- L'état partagé : + - le préchauffage en vol, son inscription, son annulation ; + - les capacités du dernier tour de chaque discussion (64 au plus). +- L'assemblage : les mêmes fonctions que le tour. + - Vue du Go, historique sans blocs ``, préambule. + - Écho, messages d'envoi, corps. + - Sentinelle « . » (rôle `builder` en mode Code), non streamé, un seul jeton. +- On s'abstient dans ces cas : + - capacités inconnues ; + - tour en cours ; + - discussion vide ; + - compaction attendue ; + - assemblage inattendu. +- Déclencheurs : + - fin de tour quand rien n'attend dans la file ; + - fin de tâche (relayée par le Go) ; + - changement de discussion ou de projet : le préchauffage en cours est + annulé ; avec `full`, la nouvelle discussion est préparée après 3 s. +- Le tour : + - `TurnHost::prewarm_keeps` compare les empreintes des messages envoyés + et des champs qui façonnent le rendu (`prewarmOptsHash`) ; + - une requête qui ne prolonge pas le préchauffage le coupe tout de suite, + sans attendre l'aller-retour par le Go. +- Télémétrie : le préchauffage est rangé dans l'anneau (`kind=prewarm`). + +## Tests + +- `chat/prewarm` : + - mode (jamais pour un preset externe) ; + - empreinte des options (l'échantillonnage n'y entre pas, un champ absent + n'est pas un champ vide) ; + - rôle de la sentinelle ; + - prolongement ; + - état partagé (un seul en vol, garde, annulation, fin) ; + - capacités bornées. +- La boucle du tour : le drapeau transmis à l'envoi, jamais vers une API + externe. +- Le vrai hôte (`tests/lot6.rs`) : + - fin de tour : requête préchauffée (sentinelle, un seul jeton, même + préfixe et mêmes outils que le tour), inscrite au Go, pas de ticket de + la façade ; + - le tour suivant la prolonge et la garde (`prewarm_keep`) ; + - fin rendue avec total et cache, télémétrie ; + - refus du Go : rien ne part ; + - annulation relayée : connexion coupée, place rendue. +- Le Go (`facade_prewarm_test.go`) : + - contrôle : clé, slots, minuit, travail annexe, moteur muet ; + - inscription et deuxième préchauffage refusé ; + - gardé par le tour qui le prolonge, annulé et relayé sinon ; + - fin unique ; + - requête déjà en vol ; + - déclencheur après tâche relayé. + +## Écarts assumés + +- L'annulation d'un préchauffage pas encore inscrit au Go (aller-retour en + cours) annule celui qui part, quel que soit son numéro : il n'y en a + qu'un. diff --git a/internal/loki/chat_prewarm.go b/internal/loki/chat_prewarm.go index 0f96093..fc6884e 100644 --- a/internal/loki/chat_prewarm.go +++ b/internal/loki/chat_prewarm.go @@ -267,10 +267,16 @@ var ( ) // prewarmKick : un déclencheur. Sans la clé, rien — ni goroutine, ni requête. +// Quand la façade tient la conversation, le déclencheur lui part +// (facade_prewarm.go) : c'est elle qui a l'historique. func (c *Conversation) prewarmKick(why string) { if prewarmMode(ReadConfig()) == "" { return } + if facadeChat() { + go func() { _ = facadePost("/internal/facade/v1/chat/prewarm", map[string]any{"why": why}) }() + return + } if prewarmAsync { go c.prewarm(why) return diff --git a/internal/loki/facade_chat.go b/internal/loki/facade_chat.go index 0e2cdeb..236802c 100644 --- a/internal/loki/facade_chat.go +++ b/internal/loki/facade_chat.go @@ -48,18 +48,18 @@ func facadeChat() bool { // facadeChatBlocker : le réglage qui garde la conversation ici, "" s'il n'y // en a pas. Ces clés opt-in demandent une machinerie restée au Go (bloc projet -// figé, préchauffage). NUDGE_IN_TOOL n'en est plus (lot 23) : la façade a sa +// figé). NUDGE_IN_TOOL n'en est plus (lot 23) : la façade a sa // sonde de gabarit ; ni les motifs de relance personnalisés (retry_patterns, // lot 30) : la façade les lit et les compile ; ni SIDE_SLOT et SLOT_PERSIST // (lot 31) : la façade demande son slot avant chaque requête // (facade_slots.go) ; ni KEEP_TURN_IMAGES (lot 32) : la façade tient l'état -// du tour, le Go range les images (facade_keepimg.go). +// du tour, le Go range les images (facade_keepimg.go) ; ni PREWARM (lot 33) : +// la façade assemble le préchauffage, le Go l'inscrit et l'annule +// (facade_prewarm.go). func facadeChatBlocker(cfg map[string]string) string { switch { case projSnapEnabled(cfg): return "PROJ_SNAPSHOT" - case prewarmMode(cfg) != "": - return "PREWARM" } return "" } diff --git a/internal/loki/facade_chat_test.go b/internal/loki/facade_chat_test.go index e9df015..9c3276b 100644 --- a/internal/loki/facade_chat_test.go +++ b/internal/loki/facade_chat_test.go @@ -277,7 +277,7 @@ func TestFacadeReglageQuiGardeLaConversation(t *testing.T) { t.Fatalf("détenteur : %d %s", w.Code, w.Body.String()) } // Posé après la décision : sans effet avant le redémarrage. - if err := SetConfigKey("PREWARM", "on"); err != nil { + if err := SetConfigKey("PROJ_SNAPSHOT", "on"); err != nil { t.Fatal(err) } if !facadeChat() { @@ -285,10 +285,10 @@ func TestFacadeReglageQuiGardeLaConversation(t *testing.T) { } oublierDecisionFacade(t) if facadeChat() { - t.Fatal("PREWARM posé : la conversation doit rester au Go") + t.Fatal("PROJ_SNAPSHOT posé : la conversation doit rester au Go") } w := appelOAIInterne(t, "GET", "/internal/v1/chat/holder", "") - if !strings.Contains(w.Body.String(), `"holder":"go"`) || !strings.Contains(w.Body.String(), "PREWARM") { + if !strings.Contains(w.Body.String(), `"holder":"go"`) || !strings.Contains(w.Body.String(), "PROJ_SNAPSHOT") { t.Fatalf("détenteur : %s", w.Body.String()) } local := func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "go") } @@ -299,7 +299,7 @@ func TestFacadeReglageQuiGardeLaConversation(t *testing.T) { } for k, want := range map[string]string{ "PROJ_SNAPSHOT": "PROJ_SNAPSHOT", "NUDGE_IN_TOOL": "", "SIDE_SLOT": "", - "SLOT_PERSIST": "", "KEEP_TURN_IMAGES": "", "PREWARM": "PREWARM", + "SLOT_PERSIST": "", "KEEP_TURN_IMAGES": "", "PREWARM": "", } { if got := facadeChatBlocker(map[string]string{k: "on"}); got != want { t.Errorf("%s : %q", k, got) diff --git a/internal/loki/facade_prewarm.go b/internal/loki/facade_prewarm.go new file mode 100644 index 0000000..307c605 --- /dev/null +++ b/internal/loki/facade_prewarm.go @@ -0,0 +1,150 @@ +package loki + +import ( + "net/http" + "sync" +) + +// facade_prewarm.go — PREWARM quand la façade tient la conversation (lot 33 de +// la migration). La façade assemble la requête du prochain tour (elle tient +// l'historique) et l'envoie ; le Go garde ce qui décide si elle peut partir +// et ce qui l'annule : +// +// - /internal/v1/prewarm/check : le moteur est-il libre d'un vrai travail +// (travail annexe, bench, optimiseur), loin de minuit, en bonne santé, et +// sur quel slot viser (un seul slot, ou celui de la discussion avec +// SIDE_SLOT) ; +// - /internal/v1/prewarm/begin : recharge l'état SLOT_PERSIST de la +// discussion, puis inscrit le préchauffage comme LA requête en vol +// (prewarmBegin). Toute requête de Loki qui part ensuite l'annule — sauf +// un tour de la façade qui le prolonge (/internal/v1/oai/begin, +// prewarm_keep) : l'annulation est relayée à la façade, qui coupe sa +// connexion au moteur ; +// - /internal/v1/prewarm/end : la fin, avec ce que le moteur a repris du +// cache (noteEnginePrompt) ; +// - après une tâche planifiée, le déclencheur part à la façade +// (prewarmKick). + +// Préchauffages de la façade inscrits ici : id → fin. +var facadePrewarms = struct { + mu sync.Mutex + next uint64 + m map[uint64]func() +}{m: map[uint64]func(){}} + +// handleInternalPrewarmCheck : POST {} → {"ok":true,"slot":-1|0} ou +// {"ok":false,"why":"…"}. +func handleInternalPrewarmCheck(w http.ResponseWriter, r *http.Request) { + var in struct{} + if !readInternalJSON(w, r, &in) { + return + } + refuse := func(why string) { + sendJSON(w, http.StatusOK, map[string]any{"ok": false, "why": why}) + } + if prewarmMode(ReadConfig()) == "" { + refuse("clé absente") + return + } + ep := resolveChatEndpoint() + if ep.External { + refuse("preset externe") + return + } + if sideJobs.Load() > 0 || benchRunning() || tuneGuard() != nil { + refuse("moteur occupé") + return + } + if prewarmNearMidnight(prewarmNow()) { + refuse("minuit proche") + return + } + if !healthCheck() { + refuse("moteur indisponible") + return + } + // Un seul slot, ou les deux de SIDE_SLOT : le préchauffage vise alors le + // slot de la discussion (id_slot 0), celui où partira le vrai tour. + slot := -1 + if n := prewarmSlots(ep); n != 1 { + if slot = engineSlotFor(ep, true); n != 2 || slot != 0 { + refuse("plusieurs slots") + return + } + } + sendJSON(w, http.StatusOK, map[string]any{"ok": true, "slot": slot}) +} + +// handleInternalPrewarmBegin : POST {"conv","n","prefix","opts"} → +// {"ok":true,"id":n} ou {"ok":false} (une requête de Loki est en vol, ou un +// autre préchauffage). n : messages du préfixe préchauffé (sans la +// sentinelle) ; prefix : son empreinte cumulée (perfPrefix) ; opts : celle +// des champs qui façonnent le rendu. +func handleInternalPrewarmBegin(w http.ResponseWriter, r *http.Request) { + var in struct { + Conv string `json:"conv"` + N int `json:"n"` + Prefix uint64 `json:"prefix"` + Opts uint64 `json:"opts"` + } + if !readInternalJSON(w, r, &in) { + return + } + // SLOT_PERSIST : l'état gardé de la discussion d'abord — le préchauffage + // le prolonge au lieu de tout recalculer. + slotPersistRestore(resolveChatEndpoint(), in.Conv) + facadePrewarms.mu.Lock() + facadePrewarms.next++ + id := facadePrewarms.next + facadePrewarms.mu.Unlock() + run := &prewarmRun{n: in.N, prefix: in.Prefix, opts: in.Opts, done: make(chan struct{})} + // Annulé par une autre requête : la façade coupe sa connexion. + run.cancel = func() { + go func() { _ = facadePost("/internal/facade/v1/chat/prewarm/cancel", map[string]any{"id": id}) }() + } + end, ok := prewarmBegin(run) + if !ok { + sendJSON(w, http.StatusOK, map[string]any{"ok": false}) + return + } + facadePrewarms.mu.Lock() + facadePrewarms.m[id] = func() { + end() + close(run.done) + } + facadePrewarms.mu.Unlock() + sendJSON(w, http.StatusOK, map[string]any{"ok": true, "id": id}) +} + +// handleInternalPrewarmEnd : POST {"id","total","cached"} → 204. total : la +// taille du prompt préchauffé (0 : inconnue) ; cached : ce que le moteur a +// repris de son cache (absent : inconnu). +func handleInternalPrewarmEnd(w http.ResponseWriter, r *http.Request) { + var in struct { + ID uint64 `json:"id"` + Total int `json:"total"` + Cached *int `json:"cached"` + } + if !readInternalJSON(w, r, &in) { + return + } + facadePrewarms.mu.Lock() + end := facadePrewarms.m[in.ID] + delete(facadePrewarms.m, in.ID) + facadePrewarms.mu.Unlock() + if end == nil { + sendJSON(w, http.StatusNotFound, map[string]any{"error": "préchauffage inconnu"}) + return + } + end() + // Après une tâche, c'est lui qui retrouve (ou non) la conversation dans le + // cache RAM : la vérification de llm_slots.go se fait ici. + if in.Total > 0 { + cached := 0 + if in.Cached != nil { + cached = *in.Cached + } + noteEnginePrompt(in.Total, cached, in.Cached != nil) + } + w.WriteHeader(http.StatusNoContent) +} diff --git a/internal/loki/facade_prewarm_test.go b/internal/loki/facade_prewarm_test.go new file mode 100644 index 0000000..b621c95 --- /dev/null +++ b/internal/loki/facade_prewarm_test.go @@ -0,0 +1,167 @@ +package loki + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" +) + +// moteurSain : un llama-server de test qui répond à /health. +func moteurSain(t *testing.T) { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + u, _ := url.Parse(srv.URL) + if err := SetConfigKey("PORT", u.Port()); err != nil { + t.Fatal(err) + } +} + +func prewarmPost(t *testing.T, path, body string) (int, map[string]any) { + t.Helper() + w := appelOAIInterne(t, "POST", path, body) + var out map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &out) + return w.Code, out +} + +// Le moteur est-il libre d'un vrai travail, et sur quel slot viser ? +func TestFacadePrewarmCheck(t *testing.T) { + testHome(t) + t.Setenv(internalTokenEnv, jetonTest) + check := func() map[string]any { + _, out := prewarmPost(t, "/internal/v1/prewarm/check", `{}`) + return out + } + if out := check(); out["ok"] != false || out["why"] != "clé absente" { + t.Fatalf("sans la clé : %v", out) + } + if err := SetConfigKey("PREWARM", "on"); err != nil { + t.Fatal(err) + } + moteurSain(t) + prevSlots, prevNow := prewarmSlots, prewarmNow + t.Cleanup(func() { prewarmSlots, prewarmNow = prevSlots, prevNow }) + prewarmSlots = func(chatEndpoint) int { return 1 } + prewarmNow = func() time.Time { return time.Date(2026, 10, 16, 14, 0, 0, 0, time.Local) } + if out := check(); out["ok"] != true || out["slot"] != float64(-1) { + t.Fatalf("un seul slot : %v", out) + } + prewarmSlots = func(chatEndpoint) int { return 4 } + if out := check(); out["ok"] != false || out["why"] != "plusieurs slots" { + t.Fatalf("quatre slots : %v", out) + } + prewarmSlots = func(chatEndpoint) int { return 1 } + prewarmNow = func() time.Time { return time.Date(2026, 10, 16, 23, 55, 0, 0, time.Local) } + if out := check(); out["why"] != "minuit proche" { + t.Fatalf("minuit : %v", out) + } + prewarmNow = func() time.Time { return time.Date(2026, 10, 16, 14, 0, 0, 0, time.Local) } + sideJobs.Add(1) + out := check() + sideJobs.Add(-1) + if out["why"] != "moteur occupé" { + t.Fatalf("travail annexe : %v", out) + } + if err := SetConfigKey("PORT", "1"); err != nil { + t.Fatal(err) + } + if out := check(); out["why"] != "moteur indisponible" { + t.Fatalf("moteur muet : %v", out) + } +} + +// Inscrit comme LA requête en vol ; annulé par toute autre requête (la +// façade l'apprend), sauf un tour qui le prolonge ; rendu à la fin. +func TestFacadePrewarmEnVol(t *testing.T) { + testHome(t) + f := demarrerFacade(t) + begin := func() (bool, float64) { + _, out := prewarmPost(t, "/internal/v1/prewarm/begin", `{"conv":"c1","n":3,"prefix":11,"opts":22}`) + id, _ := out["id"].(float64) + return out["ok"] == true, id + } + ok, id := begin() + if !ok || id == 0 || !prewarmLive.Load() || enVol() != 1 { + t.Fatalf("inscription : ok=%v id=%v en vol=%d", ok, id, enVol()) + } + if ok, _ := begin(); ok { + t.Fatal("deux préchauffages en vol") + } + // Un tour de la façade qui le prolonge : gardé. + w := appelOAIInterne(t, "POST", "/internal/v1/oai/begin", `{"main_conv":"c1","prewarm_keep":true}`) + var tk struct { + Ticket uint64 `json:"ticket"` + } + _ = json.Unmarshal(w.Body.Bytes(), &tk) + if !prewarmLive.Load() { + t.Fatal("préchauffage annulé par le tour qui le prolonge") + } + appelOAIInterne(t, "POST", "/internal/v1/oai/end", fmt.Sprintf(`{"ticket":%d}`, tk.Ticket)) + // Toute autre requête : annulé, la façade coupe sa connexion. + w = appelOAIInterne(t, "POST", "/internal/v1/oai/begin", `{"main_conv":"c1"}`) + _ = json.Unmarshal(w.Body.Bytes(), &tk) + if prewarmLive.Load() { + t.Fatal("préchauffage gardé malgré une autre requête") + } + deadline := time.Now().Add(2 * time.Second) + want := fmt.Sprintf(`POST /internal/facade/v1/chat/prewarm/cancel jeton {"id":%d}`, int(id)) + for !recuContient(f.recu(), want) { + if time.Now().After(deadline) { + t.Fatalf("annulation non relayée : %v", f.recu()) + } + time.Sleep(10 * time.Millisecond) + } + appelOAIInterne(t, "POST", "/internal/v1/oai/end", fmt.Sprintf(`{"ticket":%d}`, tk.Ticket)) + // La fin rend la place, une seule fois. + if code, _ := prewarmPost(t, "/internal/v1/prewarm/end", fmt.Sprintf(`{"id":%d,"total":900,"cached":800}`, int(id))); code != 204 { + t.Fatalf("fin : %d", code) + } + if enVol() != 0 { + t.Fatalf("encore %d requête(s) en vol", enVol()) + } + if code, _ := prewarmPost(t, "/internal/v1/prewarm/end", fmt.Sprintf(`{"id":%d}`, int(id))); code != 404 { + t.Fatalf("fin répétée : %d", code) + } + // Une requête déjà en vol : rien ne part. + w = appelOAIInterne(t, "POST", "/internal/v1/oai/begin", `{"main_conv":"c1"}`) + _ = json.Unmarshal(w.Body.Bytes(), &tk) + if ok, _ := begin(); ok { + t.Fatal("préchauffage inscrit malgré une requête en vol") + } + appelOAIInterne(t, "POST", "/internal/v1/oai/end", fmt.Sprintf(`{"ticket":%d}`, tk.Ticket)) +} + +// Après une tâche planifiée, le déclencheur part à la façade. +func TestFacadePrewarmApresTache(t *testing.T) { + testHome(t) + f := demarrerFacade(t) + if err := SetConfigKey("PREWARM", "on"); err != nil { + t.Fatal(err) + } + conv.prewarmKick(prewarmAfterTask) + deadline := time.Now().Add(2 * time.Second) + want := `POST /internal/facade/v1/chat/prewarm jeton {"why":"après-tâche"}` + for !recuContient(f.recu(), want) { + if time.Now().After(deadline) { + t.Fatalf("déclencheur non relayé : %v", f.recu()) + } + time.Sleep(10 * time.Millisecond) + } +} + +func recuContient(l []string, s string) bool { + for _, x := range l { + if strings.TrimSpace(x) == s { + return true + } + } + return false +} diff --git a/internal/loki/web_internal.go b/internal/loki/web_internal.go index 6a36e7d..e5b63b1 100644 --- a/internal/loki/web_internal.go +++ b/internal/loki/web_internal.go @@ -75,6 +75,12 @@ func withInternalAPI(next http.Handler) http.Handler { handleInternalTurnPerf(w, r) case "/internal/v1/turn/slot": handleInternalTurnSlot(w, r) + case "/internal/v1/prewarm/check": + handleInternalPrewarmCheck(w, r) + case "/internal/v1/prewarm/begin": + handleInternalPrewarmBegin(w, r) + case "/internal/v1/prewarm/end": + handleInternalPrewarmEnd(w, r) case "/internal/v1/turn/keepimg": handleInternalTurnKeepImg(w, r) case "/internal/v1/turn/keepimg/store": diff --git a/internal/loki/web_internal_oai.go b/internal/loki/web_internal_oai.go index 44bd020..086638c 100644 --- a/internal/loki/web_internal_oai.go +++ b/internal/loki/web_internal_oai.go @@ -70,6 +70,10 @@ func handleInternalOAIBegin(w http.ResponseWriter, r *http.Request) { // MainConv : étape d'un tour de cette discussion, tenue par la façade // (LOKI_CHAT=rust) — comptée comme la requête principale du Go. MainConv string `json:"main_conv"` + // PrewarmKeep : étape d'un tour de la façade qui prolonge exactement + // le préchauffage en vol (PREWARM, facade_prewarm.go) — il n'est pas + // annulé, la requête le suit dans le même slot. + PrewarmKeep bool `json:"prewarm_keep"` } if r.Method != http.MethodPost || json.NewDecoder(r.Body).Decode(&in) != nil { sendJSON(w, http.StatusBadRequest, map[string]any{"error": "POST JSON attendu"}) @@ -80,7 +84,11 @@ func handleInternalOAIBegin(w http.ResponseWriter, r *http.Request) { if in.Read { end, seq = engineRequestBeginSide(func(*prewarmRun) bool { return true }, true) } else { - end, seq = engineRequestBeginSide(nil, in.Side) + var keep func(*prewarmRun) bool + if in.PrewarmKeep { + keep = func(*prewarmRun) bool { return true } + } + end, seq = engineRequestBeginSide(keep, in.Side) } if in.PostV1 { refuse := func(retry, msg, code string) { diff --git a/rust/loki-rs/src/chat/generate.rs b/rust/loki-rs/src/chat/generate.rs index 840812d..b697394 100644 --- a/rust/loki-rs/src/chat/generate.rs +++ b/rust/loki-rs/src/chat/generate.rs @@ -22,8 +22,9 @@ //! corrections du builder) suit le tour, avant la compaction de fin //! ([`super::verify`]) ; toute compaction rend l'état des critères au résumé. //! -//! Écarts assumés (restent au Go) : PROJ_SNAPSHOT (refusé par -//! `/turn/view`), préchauffage. Les images gardées (KEEP_TURN_IMAGES, lot 32) +//! Écart assumé (reste au Go) : PROJ_SNAPSHOT (refusé par `/turn/view`). Le +//! préchauffage (PREWARM, lot 33) part en fin de tour ([`super::prewarm`]). +//! Les images gardées (KEEP_TURN_IMAGES, lot 32) //! partent en début et en fin de tour quand il le faut //! ([`ChatService::kept_images_turn_start`]). La télémétrie des complétions //! (tours, vérifications, résumés) rejoint l'anneau du Go @@ -684,7 +685,7 @@ impl ChatService { /// Changement de projet (`/api/projects`, action switch), remis par le /// Go : refusé pendant un tour. - pub async fn project_switch(&self, slug: &str) -> Result<(), Refused> { + pub async fn project_switch(self: &Arc, slug: &str) -> Result<(), Refused> { if self.sessions.conv().is_generating() { return Err(Refused { status: 409, @@ -697,7 +698,9 @@ impl ChatService { .map_err(|e| Refused { status: 400, error: e.to_string(), - }) + })?; + self.prewarm_after_switch().await; + Ok(()) } /// Suppression d'un projet (`deleteProject`) : ses discussions d'abord, @@ -752,6 +755,9 @@ impl ChatService { ); let local = !ep.external; let conv_id = conv.id(); + // Capacités demandées pour ce tour, retenues en mémoire (jamais + // persistées) pour préchauffer le suivant (PREWARM). + self.shared.prewarm.note_caps(&conv_id, prep.caps.clone()); let agent = prep.cap("Agent"); let mem_off = prep.turn_caps.get("Mem").and_then(Value::as_str) == Some("off"); let ccaps = CompactCaps { @@ -1105,6 +1111,9 @@ impl ChatService { { eprintln!("[chat] fin de tour non signalée : {e}"); } + // Plus rien ne suit : le prochain tour peut être préparé pendant + // que l'utilisateur lit (PREWARM, off par défaut). + self.prewarm_kick(super::prewarm::TURN_END).await; } } } @@ -1160,6 +1169,7 @@ impl ChatShared { images: super::host::ImageCache::new(img_dir), patterns: super::stream::RetryPatterns::builtin(), tpl: std::sync::Arc::new(super::tplprobe::TplProbe::new()), + prewarm: Default::default(), } } } diff --git a/rust/loki-rs/src/chat/host.rs b/rust/loki-rs/src/chat/host.rs index 385d1c8..c56db96 100644 --- a/rust/loki-rs/src/chat/host.rs +++ b/rust/loki-rs/src/chat/host.rs @@ -196,7 +196,7 @@ impl EngineSummarizer { { o.insert("id_slot".into(), Value::from(n)); } - _gate = engine_begin(store, "", slot == Some(1)).await; + _gate = engine_begin(store, "", slot == Some(1), false).await; } let body = serde_json::to_vec(&payload).unwrap_or_default(); let resp = self @@ -328,6 +328,8 @@ pub struct ChatShared { pub patterns: Arc, /// Sonde de gabarit (`chat_tplprobe.go`) : derniers verdicts et cache. pub tpl: Arc, + /// Préchauffage (PREWARM) : celui en vol, capacités des discussions. + pub prewarm: super::prewarm::Prewarm, } /// Ce que la compaction sait du tour. @@ -397,15 +399,21 @@ pub async fn engine_slot_for(store: &Store, conv: &str, main: bool) -> Option Option { +/// préchauffage est annulé — sauf `keep_prewarm`, une étape de tour qui le +/// prolonge. `main_conv` : étape d'un tour de cette discussion ; `side` : +/// requête sur le second slot. `None` si le Go ne répond pas — la requête +/// part quand même. +pub async fn engine_begin( + store: &Store, + main_conv: &str, + side: bool, + keep_prewarm: bool, +) -> Option { match store .internal_json( hyper::Method::POST, "/internal/v1/oai/begin", - Some(&json!({ "main_conv": main_conv, "side": side })), + Some(&json!({ "main_conv": main_conv, "side": side, "prewarm_keep": keep_prewarm })), ) .await { @@ -461,14 +469,14 @@ impl Drop for GateTicket { } impl GoTurnHost { - async fn gate_begin(&self) -> Option { + async fn gate_begin(&self, keep_prewarm: bool) -> Option { let main_conv = if self.perf_kind == super::perf::KIND_MAIN { self.conv_id.as_str() } else { "" }; let side = self.side_slot.load(Ordering::SeqCst); - engine_begin(&self.shared.store, main_conv, side).await + engine_begin(&self.shared.store, main_conv, side, keep_prewarm).await } } @@ -527,13 +535,23 @@ pub fn see_image_message(label: &str, img: Value) -> Message { } impl TurnHost for GoTurnHost { - async fn send(&self, ep: &ChatEndpoint, body: Vec) -> Result { + async fn send( + &self, + ep: &ChatEndpoint, + body: Vec, + keep_prewarm: bool, + ) -> Result { // Moteur local : requête comptée en vol jusqu'à la fin de lecture du // flux (le ticket voyage avec le corps). Rien pour une API externe. + // Un préchauffage qu'elle ne prolonge pas est coupé tout de suite (le + // Go l'annulerait aussi, un aller-retour plus tard). let gate = if ep.external { None } else { - self.gate_begin().await + if !keep_prewarm { + self.shared.prewarm.stop(); + } + self.gate_begin(keep_prewarm).await }; match self.shared.http.post(&ep.url, &ep.key, body).await { Ok(resp) => { @@ -720,6 +738,15 @@ impl TurnHost for GoTurnHost { self.shared.echo_refusals.refuse(model); } + fn prewarm_keeps(&self, prefix: &[u64], payload: &Value) -> bool { + self.perf_kind == super::perf::KIND_MAIN + && self.shared.prewarm.live() + && self + .shared + .prewarm + .keeps(prefix, super::prewarm::opts_hash(payload)) + } + async fn keep_images(&self) -> Option { // Fil principal d'une discussion seulement : la vérification ne // persiste pas son historique. diff --git a/rust/loki-rs/src/chat/mod.rs b/rust/loki-rs/src/chat/mod.rs index fbd6997..2a8a384 100644 --- a/rust/loki-rs/src/chat/mod.rs +++ b/rust/loki-rs/src/chat/mod.rs @@ -33,6 +33,7 @@ pub mod lease; pub mod nudges; pub mod perf; pub mod persist; +pub mod prewarm; pub mod sessions; pub mod stream; pub mod tplprobe; diff --git a/rust/loki-rs/src/chat/perf.rs b/rust/loki-rs/src/chat/perf.rs index 44b6b5b..7db38a9 100644 --- a/rust/loki-rs/src/chat/perf.rs +++ b/rust/loki-rs/src/chat/perf.rs @@ -19,6 +19,7 @@ use crate::store::Store; pub const KIND_MAIN: &str = "main"; pub const KIND_VERIFY: &str = "verify"; pub const KIND_COMPACT: &str = "compact"; +pub const KIND_PREWARM: &str = "prewarm"; /// `perfRecFromWire` : une réponse NON streamée (compaction), lue d'un /// bloc, au format des stats d'une complétion streamée. Jamais d'échec : un diff --git a/rust/loki-rs/src/chat/prewarm.rs b/rust/loki-rs/src/chat/prewarm.rs new file mode 100644 index 0000000..218dd93 --- /dev/null +++ b/rust/loki-rs/src/chat/prewarm.rs @@ -0,0 +1,513 @@ +//! PREWARM tenu par la façade (lot 33) — port de `chat_prewarm.go`. +//! +//! Préparer le prochain tour pendant que l'utilisateur lit. Certains +//! recalculs sont inévitables (prompt réécrit par une compaction, dernier +//! message assistant rendu autrement, discussion reprise après une tâche qui +//! a pris le slot) ; sans la clé, ils tombent sur le premier jeton du message +//! suivant. +//! +//! Avec la clé, la façade envoie au moteur local, dès qu'il est libre, la +//! requête du prochain tour telle qu'elle partira — assemblée par les MÊMES +//! fonctions que le tour — suivie d'un message utilisateur sentinelle « . », +//! avec `max_tokens` 1. Le moteur calcule tout le préfixe ; le vrai tour ne +//! calcule plus que son propre message. Rien n'est persisté. +//! +//! Le Go garde ce qui décide si le préchauffage peut partir et ce qui +//! l'annule (`facade_prewarm.go`) : il l'inscrit comme LA requête en vol, et +//! toute requête de Loki qui part ensuite l'annule — sauf un tour de la +//! façade dont la requête le PROLONGE exactement (empreintes des messages, +//! des outils et des arguments du gabarit), qui le suit dans le même slot. +//! Aucune erreur ne remonte. + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use serde_json::{Value, json}; +use tokio::sync::watch; + +use super::compact::{compact_needed, ctx_pending, strip_context_updates}; +use super::generate::ChatService; +use super::history::inject_prefix; +use super::perf::{self, wire_stats}; +use super::turn::{PerfRec, perf_prefix}; +use super::types::Message; +use super::wire::{ + ChatPayloadOpts, build_chat_payload, turn_echo_policy, turn_reasoning, wire_messages, +}; +use crate::engstatus::is_external; + +/// Le message utilisateur qui tient la place du vrai. +pub const SENTINEL: &str = "."; + +/// Déclencheurs (pour le journal). +pub const TURN_END: &str = "fin-tour"; +pub const AFTER_TASK: &str = "après-tâche"; +pub const SWITCH: &str = "discussion"; + +/// PREWARM=full : préparer la discussion ouverte si l'utilisateur y reste. +pub const SWITCH_DELAY: Duration = Duration::from_secs(3); +/// Borne d'un préchauffage (un MoE aux experts en RAM met des minutes sur 20k +/// jetons). Toute vraie requête l'annule bien avant. +pub const TIMEOUT: Duration = Duration::from_secs(15 * 60); +/// Attente de la fin d'un préchauffage remplacé. +const STOP_WAIT: Duration = Duration::from_secs(5); +/// Capacités gardées, une par discussion. +const CAPS_MAX: usize = 64; + +/// `prewarmMode` : « on », « full » ou "" (off). Jamais pour un preset externe. +pub fn prewarm_mode(cfg: &HashMap) -> &'static str { + if is_external(cfg) { + return ""; + } + match cfg + .get("PREWARM") + .map(|v| v.trim().to_lowercase()) + .as_deref() + { + Some("on" | "1" | "true" | "yes" | "oui") => "on", + Some("full") => "full", + _ => "", + } +} + +/// `prewarmOptsKeys` : les champs du corps, hors messages, qui changent le +/// prompt rendu par le gabarit (l'échantillonnage n'y touche pas). +const OPTS_KEYS: [&str; 5] = [ + "tools", + "parallel_tool_calls", + "tool_choice", + "chat_template_kwargs", + "reasoning_effort", +]; + +/// `prewarmOptsHash` : FNV-1a 64 de ces champs. Comparée seulement aux +/// empreintes calculées ici, d'un corps à l'autre. +pub fn opts_hash(payload: &Value) -> u64 { + const OFFSET: u64 = 0xcbf2_9ce4_8422_2325; + const PRIME: u64 = 0x0000_0100_0000_01b3; + let mut h = OFFSET; + let mut eat = |b: &[u8]| { + for &c in b { + h = (h ^ u64::from(c)).wrapping_mul(PRIME); + } + }; + for k in OPTS_KEYS { + eat(k.as_bytes()); + eat(&[0]); + if let Some(v) = payload.get(k) { + eat(&serde_json::to_vec(v).unwrap_or_default()); + } + eat(&[0]); + } + h +} + +/// `withTurnRole(caps, ".")` : en mode Code, le rôle d'un message ordinaire. +pub fn sentinel_caps(caps: &Value) -> Value { + let mut c = caps.clone(); + let code = c.get("Code").and_then(Value::as_bool).unwrap_or(false); + let role_empty = c + .get("Role") + .and_then(Value::as_str) + .unwrap_or("") + .is_empty(); + if code + && role_empty + && let Some(o) = c.as_object_mut() + { + o.insert("Role".into(), Value::from("builder")); + } + c +} + +/// Un préchauffage inscrit au Go : `prefix` couvre ses `n` premiers messages +/// (sans la sentinelle), `opts` les champs qui façonnent le rendu. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Run { + pub id: u64, + pub n: usize, + pub prefix: u64, + pub opts: u64, +} + +impl Run { + /// `matches` : la requête dont `prefix` est l'empreinte cumulée des + /// messages prolonge-t-elle exactement ce préchauffage ? + pub fn matches(&self, prefix: &[u64], opts: u64) -> bool { + self.n > 0 + && prefix.len() > self.n + && prefix[self.n - 1] == self.prefix + && opts == self.opts + } +} + +struct Current { + /// Rempli une fois inscrit au Go. + run: Option, + cancel: watch::Sender, + done: watch::Receiver, +} + +/// L'état partagé du préchauffage (un exemplaire par processus). +#[derive(Default)] +pub struct Prewarm { + cur: Mutex>, + /// Capacités du dernier tour de chaque discussion, en mémoire seulement. + caps: Mutex>, + /// Raisons d'échec déjà journalisées. + logged: Mutex>, +} + +impl Prewarm { + fn cur(&self) -> std::sync::MutexGuard<'_, Option> { + self.cur.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Un préchauffage est-il en vol (ou en train de partir) ? + pub fn live(&self) -> bool { + self.cur().is_some() + } + + /// La requête d'un tour le prolonge-t-elle exactement ? + pub fn keeps(&self, prefix: &[u64], opts: u64) -> bool { + self.cur() + .as_ref() + .and_then(|c| c.run) + .is_some_and(|r| r.matches(prefix, opts)) + } + + /// `prewarmStop` : annule le préchauffage en vol et rend de quoi attendre + /// sa fin (`None` : rien en vol). + pub fn stop(&self) -> Option> { + let c = self.cur().take()?; + let _ = c.cancel.send(true); + Some(c.done) + } + + /// Annulation relayée par le Go (une autre requête est partie). Un + /// préchauffage pas encore inscrit est annulé aussi : il n'y en a qu'un. + pub fn cancel_id(&self, id: u64) -> bool { + let g = self.cur(); + match g.as_ref() { + Some(c) if c.run.is_none_or(|r| r.id == id) => { + let _ = c.cancel.send(true); + true + } + _ => false, + } + } + + /// `prewarmNoteCaps`. + pub fn note_caps(&self, conv: &str, caps: Value) { + if conv.is_empty() { + return; + } + let mut m = self.caps.lock().unwrap_or_else(|e| e.into_inner()); + if !m.contains_key(conv) + && m.len() >= CAPS_MAX + && let Some(k) = m.keys().next().cloned() + { + m.remove(&k); + } + m.insert(conv.to_string(), caps); + } + + pub fn caps_for(&self, conv: &str) -> Option { + self.caps + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(conv) + .cloned() + } + + fn log_once(&self, key: &str, msg: &str) { + if self + .logged + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(key.to_string()) + { + eprintln!("[prewarm] {msg}"); + } + } + + /// Inscrit le préchauffage qui part ; `None` si un autre est déjà là. + fn register(&self) -> Option<(watch::Receiver, watch::Sender)> { + let mut g = self.cur(); + if g.is_some() { + return None; + } + let (cancel, cancelled) = watch::channel(false); + let (done_tx, done) = watch::channel(false); + *g = Some(Current { + run: None, + cancel, + done, + }); + Some((cancelled, done_tx)) + } + + fn set_run(&self, run: Run) { + if let Some(c) = self.cur().as_mut() { + c.run = Some(run); + } + } + + /// Fin du préchauffage `run` (ou de celui qui partait) : la place est + /// rendue s'il est toujours le courant. + fn release(&self, id: Option) { + let mut g = self.cur(); + if g.as_ref().is_some_and(|c| c.run.map(|r| r.id) == id) { + *g = None; + } + } +} + +/// Ce que l'assemblage rend : le corps, et l'empreinte du préfixe préchauffé. +#[derive(Debug, Clone, PartialEq)] +pub struct Body { + pub payload: Value, + pub n: usize, + pub prefix: u64, + pub opts: u64, +} + +impl ChatService { + /// `prewarmKick` : un déclencheur. Sans la clé, rien — ni tâche, ni + /// requête. + pub async fn prewarm_kick(self: &Arc, why: &'static str) { + let cfg = self.shared.store.config().await; + if prewarm_mode(&cfg).is_empty() { + return; + } + let me = self.clone(); + tokio::spawn(async move { me.prewarm(why).await }); + } + + /// `prewarmAfterSwitch` : changement de discussion. Le préchauffage de + /// l'ancienne est annulé ; avec PREWARM=full, la nouvelle est préparée si + /// l'utilisateur y reste 3 s. + pub async fn prewarm_after_switch(self: &Arc) { + self.shared.prewarm.stop(); // sa fin n'est pas attendue + let cfg = self.shared.store.config().await; + if prewarm_mode(&cfg) != "full" { + return; + } + let id = self.sessions.conv().id(); + let me = self.clone(); + tokio::spawn(async move { + tokio::time::sleep(SWITCH_DELAY).await; + if me.sessions.conv().id() == id { + me.prewarm_kick(SWITCH).await; + } + }); + } + + /// `prewarmBodyFor` : le corps du préchauffage de la discussion ouverte, + /// ou la raison de s'abstenir. Lecture seule. + pub async fn prewarm_body(&self) -> Result { + let env = self.turn_env().await; + let conv_id = self.sessions.conv().id(); + let caps = self + .shared + .prewarm + .caps_for(&conv_id) + .ok_or("capacités du dernier tour inconnues")?; + let turn_caps = sentinel_caps(&caps); + let conv = self.sessions.conv(); + if conv.is_generating() { + return Err("moteur occupé"); + } + let (mut msgs, ctx_now, peak) = conv.with_state(|st| { + let msgs = st.messages.clone(); + let used = ctx_pending(st.ctx_used, st.ctx_used_len, &msgs, &env.echo); + (msgs, used, st.gen_peak) + }); + // Sans message utilisateur, le contexte projet irait en tête de la + // sentinelle : rien à préparer. + if !msgs.iter().any(|m| m.role == "user") { + return Err("discussion vide"); + } + msgs.push(Message { + role: "user".into(), + content: Some(Value::from(SENTINEL)), + ..Default::default() + }); + // Le prochain tour compactera d'abord : son prompt ne sera pas celui-ci. + if compact_needed(&env.ccfg, &msgs, ctx_now, peak) { + return Err("compaction attendue au prochain tour"); + } + let (head, prefix, tools) = self + .turn_view(&turn_caps) + .await + .map_err(|_| "vue du tour refusée")?; + let mut sent = head; + sent.extend(strip_context_updates(msgs)); + let sent = inject_prefix(sent, &prefix); + let reasoning = turn_reasoning(&env.cfg, &self.shared.effort); + let pol = turn_echo_policy(false, || env.echo.clone()); + self.shared.images.prefetch(&self.shared.store, &sent).await; + let (wire, _) = wire_messages(&sent, false, &pol); + let n = wire.len(); + let last_ok = wire.last().is_some_and(|m| { + m.role == "user" && m.content.as_ref().and_then(Value::as_str) == Some(SENTINEL) + }); + if n < 2 || !last_ok { + return Err("assemblage inattendu"); + } + let mut payload = build_chat_payload( + &env.ep, + &wire, + 0.7, + &ChatPayloadOpts { + tools: &tools, + disable_tools: false, + tool_choice_none: false, + effort: &reasoning.effort, + kwargs: reasoning.kwargs.as_ref(), + no_sampling: false, + }, + &env.cfg, + &|r| self.shared.images.get(r), + ); + // Seuls les champs de réponse changent : un seul jeton, non streamé. + if let Some(o) = payload.as_object_mut() { + o.insert("stream".into(), Value::Bool(false)); + o.remove("stream_options"); + o.insert("max_tokens".into(), json!(1)); + } + let opts = opts_hash(&payload); + Ok(Body { + payload, + n: n - 1, + prefix: perf_prefix(&wire)[n - 2], + opts, + }) + } + + /// `prewarm` : un préchauffage complet, ou rien. Jamais d'erreur rendue. + pub async fn prewarm(self: &Arc, why: &str) { + let cfg = self.shared.store.config().await; + if prewarm_mode(&cfg).is_empty() { + return; + } + let env = self.turn_env().await; + if env.ep.external { + return; + } + // Un seul en vol : le nouveau déclencheur remplace l'ancien. + if let Some(mut done) = self.shared.prewarm.stop() { + let ended = tokio::time::timeout(STOP_WAIT, done.wait_for(|d| *d)).await; + if !matches!(ended, Ok(Ok(_))) { + return; + } + } + let store = &self.shared.store; + let slot = match store + .internal_post("/internal/v1/prewarm/check", &json!({})) + .await + { + Ok((code, v)) if code.is_success() && v["ok"] == true => { + v["slot"].as_i64().unwrap_or(-1) + } + _ => return, + }; + let Ok(mut body) = self.prewarm_body().await else { + return; + }; + if slot >= 0 + && let Some(o) = body.payload.as_object_mut() + { + o.insert("id_slot".into(), json!(slot)); + } + let Ok(raw_body) = serde_json::to_vec(&body.payload) else { + return; + }; + let Some((mut cancelled, done)) = self.shared.prewarm.register() else { + return; + }; + let conv_id = self.sessions.conv().id(); + let begun = store + .internal_post( + "/internal/v1/prewarm/begin", + &json!({ "conv": conv_id, "n": body.n, "prefix": body.prefix, "opts": body.opts }), + ) + .await; + let id = match begun { + Ok((code, v)) if code.is_success() && v["ok"] == true => v["id"].as_u64(), + _ => None, + }; + let Some(id) = id else { + // Une requête de Loki est partie entre-temps. + self.shared.prewarm.release(None); + let _ = done.send(true); + return; + }; + self.shared.prewarm.set_run(Run { + id, + n: body.n, + prefix: body.prefix, + opts: body.opts, + }); + let (mut total, mut cached) = (0i64, None::); + if !*cancelled.borrow() { + let send = async { + let resp = self + .shared + .http + .post(&env.ep.url, &env.ep.key, raw_body) + .await; + match resp { + Err(e) => Err(format!("moteur injoignable ({why}) : {e}")), + Ok(r) => { + let status = r.status(); + let raw = r.bytes().await.map_err(|e| e.to_string())?; + let raw = &raw[..raw.len().min(1 << 20)]; + Ok((status, raw.to_vec())) + } + } + }; + let out: Option), String>> = tokio::select! { + o = tokio::time::timeout(TIMEOUT, send) => o.ok(), + _ = cancelled.wait_for(|c| *c) => None, + }; + match out { + Some(Err(e)) => self.shared.prewarm.log_once("net", &e), + Some(Ok((status, raw))) if status.as_u16() != 200 => { + let txt = String::from_utf8_lossy(&raw[..raw.len().min(200)]) + .trim() + .to_string(); + self.shared.prewarm.log_once( + status.as_str(), + &format!("refusé par le moteur ({why}) : {status} {txt}"), + ); + } + Some(Ok((_, raw))) => { + // Télémétrie seulement (kind=prewarm) : la réponse est jetée. + let (complete, stats) = wire_stats(&raw); + total = stats.prompt_tokens_total; + cached = stats.cache_tokens; + let rec = PerfRec { + complete, + stats, + ..Default::default() + }; + perf::record(store, perf::KIND_PREWARM, &conv_id, &rec).await; + } + None => {} + } + } + let _ = store + .internal_post( + "/internal/v1/prewarm/end", + &json!({ "id": id, "total": total, "cached": cached }), + ) + .await; + self.shared.prewarm.release(Some(id)); + let _ = done.send(true); + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/loki-rs/src/chat/prewarm/tests.rs b/rust/loki-rs/src/chat/prewarm/tests.rs new file mode 100644 index 0000000..258a190 --- /dev/null +++ b/rust/loki-rs/src/chat/prewarm/tests.rs @@ -0,0 +1,167 @@ +//! Tests du préchauffage (cas de `chat_prewarm_test.go`). + +use super::*; +use serde_json::json; + +fn cfg(kv: &[(&str, &str)]) -> HashMap { + kv.iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() +} + +fn um(s: &str) -> Message { + Message { + role: "user".into(), + content: Some(json!(s)), + ..Default::default() + } +} + +#[test] +fn mode() { + for (v, want) in [ + ("on", "on"), + (" OUI ", "on"), + ("1", "on"), + ("full", "full"), + ("FULL", "full"), + ("off", ""), + ("", ""), + ] { + assert_eq!(prewarm_mode(&cfg(&[("PREWARM", v)])), want, "{v:?}"); + } + assert_eq!(prewarm_mode(&cfg(&[])), ""); + // Jamais pour un preset externe. + assert_eq!( + prewarm_mode(&cfg(&[ + ("PREWARM", "full"), + (crate::engstatus::EXT_KEY_FLAG, "1") + ])), + "" + ); +} + +/// Seuls les champs qui changent le prompt rendu comptent ; l'échantillonnage +/// et les champs de réponse, non. +#[test] +fn empreinte_des_options() { + let base = json!({ + "model": "m", "messages": [], "stream": true, "temperature": 0.7, + "tools": [{"type": "function", "function": {"name": "bash"}}], + "tool_choice": "auto", "chat_template_kwargs": {"enable_thinking": true} + }); + let h = opts_hash(&base); + let mut same = base.clone(); + same["temperature"] = json!(0.2); + same["stream"] = json!(false); + same["max_tokens"] = json!(1); + same["id_slot"] = json!(0); + same["messages"] = json!([{"role": "user", "content": "x"}]); + assert_eq!(opts_hash(&same), h); + for (k, v) in [ + ("tools", json!([])), + ("tool_choice", json!("none")), + ("chat_template_kwargs", json!({"enable_thinking": false})), + ("reasoning_effort", json!("low")), + ("parallel_tool_calls", json!(true)), + ] { + let mut other = base.clone(); + other[k] = v; + assert_ne!(opts_hash(&other), h, "{k}"); + } + // Un champ absent n'est pas un champ vide. + let mut absent = base.clone(); + absent.as_object_mut().unwrap().remove("tool_choice"); + let mut empty = base.clone(); + empty["tool_choice"] = json!(""); + assert_ne!(opts_hash(&absent), opts_hash(&empty)); +} + +#[test] +fn role_de_la_sentinelle() { + assert_eq!( + sentinel_caps(&json!({"Code": true, "Role": ""}))["Role"], + "builder" + ); + assert_eq!(sentinel_caps(&json!({"Code": true}))["Role"], "builder"); + assert_eq!( + sentinel_caps(&json!({"Code": true, "Role": "verifier"}))["Role"], + "verifier" + ); + assert!(sentinel_caps(&json!({"Agent": true})).get("Role").is_none()); +} + +/// Une requête prolonge le préchauffage quand ses n premiers messages ont la +/// même empreinte, qu'elle en a au moins un de plus, et mêmes options. +#[test] +fn prolongement() { + let wire = vec![um("a"), um("b"), um(SENTINEL)]; + let pre = perf_prefix(&wire); + let run = Run { + id: 1, + n: 2, + prefix: pre[1], + opts: 7, + }; + let real = perf_prefix(&[um("a"), um("b"), um("vraie question")]); + assert!(run.matches(&real, 7)); + assert!(!run.matches(&real, 8), "autres options"); + assert!(!run.matches(&real[..2], 7), "rien de plus que le préfixe"); + let other = perf_prefix(&[um("a"), um("c"), um("vraie question")]); + assert!(!run.matches(&other, 7), "préfixe différent"); + let longer = perf_prefix(&[um("a"), um("b"), um("x"), um("y")]); + assert!(run.matches(&longer, 7)); + assert!(!Run { n: 0, ..run }.matches(&real, 7)); +} + +/// Un seul en vol ; gardé par la requête qui le prolonge ; annulé (par le +/// Go ou par une autre requête) ; la place rendue à sa fin seulement. +#[test] +fn etat_partage() { + let p = Prewarm::default(); + assert!(!p.live()); + let (cancelled, done) = p.register().expect("libre"); + assert!(p.live()); + assert!(p.register().is_none(), "deux en vol"); + // Pas encore inscrit au Go : rien à garder, mais annulable. + assert!(!p.keeps(&[1, 2, 3], 7)); + p.set_run(Run { + id: 4, + n: 2, + prefix: 2, + opts: 7, + }); + assert!(p.keeps(&[1, 2, 3], 7)); + assert!(!p.keeps(&[1, 2, 3], 8)); + assert!(!p.cancel_id(9), "annulation d'un autre préchauffage"); + assert!(!*cancelled.borrow()); + assert!(p.cancel_id(4)); + assert!(*cancelled.borrow()); + assert!(p.live(), "la place reste prise jusqu'à la fin"); + p.release(Some(9)); + assert!(p.live(), "fin d'un autre préchauffage"); + p.release(Some(4)); + assert!(!p.live()); + drop(done); + // stop : annule et rend de quoi attendre la fin. + let (cancelled, done) = p.register().unwrap(); + let mut wait = p.stop().expect("en vol"); + assert!(*cancelled.borrow() && !p.live()); + let _ = done.send(true); + assert!(*wait.borrow_and_update()); + assert!(p.stop().is_none()); +} + +#[test] +fn capacites_par_discussion() { + let p = Prewarm::default(); + p.note_caps("", json!({"Agent": true})); + assert!(p.caps_for("").is_none()); + p.note_caps("c1", json!({"Agent": true})); + p.note_caps("c1", json!({"Agent": false})); + assert_eq!(p.caps_for("c1"), Some(json!({"Agent": false}))); + for i in 0..CAPS_MAX + 10 { + p.note_caps(&format!("x{i}"), json!({})); + } + assert_eq!(p.caps.lock().unwrap().len(), CAPS_MAX); +} diff --git a/rust/loki-rs/src/chat/turn.rs b/rust/loki-rs/src/chat/turn.rs index 040f0cd..027bcbe 100644 --- a/rust/loki-rs/src/chat/turn.rs +++ b/rust/loki-rs/src/chat/turn.rs @@ -38,9 +38,12 @@ //! est gardée à l'essai jusqu'à la mesure de son coût par le moteur //! ([`super::keepimg`]) ; les images gardées partent avant tout résumé. //! -//! Non portés (clés opt-in, ou machinerie du moteur local restée au Go) : -//! PROJ_SNAPSHOT, préchauffage. Avec l'une de ces clés, le Go garde la -//! conversation (`facadeChatBlocker`, décidé au démarrage). +//! Préchauffage (PREWARM, lot 33) : une requête qui prolonge exactement le +//! préchauffage en vol le garde ([`TurnHost::prewarm_keeps`]) ; toute autre +//! l'annule ([`super::prewarm`]). +//! +//! Non porté : PROJ_SNAPSHOT. Avec cette clé, le Go garde la conversation +//! (`facadeChatBlocker`, décidé au démarrage). use std::collections::HashMap; use std::future::Future; @@ -138,12 +141,20 @@ pub enum CompactPhase { /// Les effets d'un tour, fournis par l'appelant. pub trait TurnHost: Send + Sync { - /// POST du corps vers l'endpoint (sa clé à lui). + /// POST du corps vers l'endpoint (sa clé à lui). `keep_prewarm` : la + /// requête prolonge exactement le préchauffage en vol, qui n'est pas + /// annulé ([`Self::prewarm_keeps`]). fn send( &self, ep: &ChatEndpoint, body: Vec, + keep_prewarm: bool, ) -> impl Future> + Send; + /// PREWARM : la requête qui part (empreintes cumulées de ses messages, + /// corps) prolonge-t-elle exactement le préchauffage en vol ? + fn prewarm_keeps(&self, _prefix: &[u64], _payload: &Value) -> bool { + false + } /// Exécute un appel d'outil (`execTool`). fn exec_tool( &self, @@ -541,7 +552,7 @@ async fn run_turn_loop( let pol = turn_echo_policy(ep.external || s.echo_off_turn, || tc.echo.clone()); host.prepare_images(&messages).await; let slot = host.engine_slot(ep).await; - let (body, echo_sent, prefix) = { + let (body, echo_sent, prefix, keep_prewarm) = { let (sent, echo_sent) = wire_messages(&messages, s.tool_choice_none, &pol); let mut payload = build_chat_payload( ep, @@ -564,10 +575,13 @@ async fn run_turn_loop( { o.insert("id_slot".into(), Value::from(n)); } + let prefix = perf_prefix(&sent); + let keep = !ep.external && host.prewarm_keeps(&prefix, &payload); ( serde_json::to_vec(&payload).unwrap_or_default(), echo_sent, - perf_prefix(&sent), + prefix, + keep, ) }; let t_req = Instant::now(); @@ -575,7 +589,7 @@ async fn run_turn_loop( if let Some(k) = kt.as_mut() { k.sending(messages.len()); } - let reply = match cancel.or_stop(host.send(ep, body)).await { + let reply = match cancel.or_stop(host.send(ep, body, keep_prewarm)).await { None => return TurnResult { extra, err: None }, Some(Ok(r)) => r, Some(Err(e)) => { diff --git a/rust/loki-rs/src/chat/turn/tests.rs b/rust/loki-rs/src/chat/turn/tests.rs index d5750e8..31e4981 100644 --- a/rust/loki-rs/src/chat/turn/tests.rs +++ b/rust/loki-rs/src/chat/turn/tests.rs @@ -48,6 +48,10 @@ struct Fake { relay: Mutex>, keep_stores: Mutex, keep_refuse: Mutex, + /// PREWARM : la requête prolonge-t-elle le préchauffage ? ce que `send` + /// a reçu. + prewarm_match: Mutex, + sent_keeps: Mutex>, } impl Fake { @@ -77,7 +81,13 @@ fn stream_of(lines: Vec, cut: bool) -> BodyStream { } impl TurnHost for Fake { - async fn send(&self, _ep: &ChatEndpoint, body: Vec) -> Result { + async fn send( + &self, + _ep: &ChatEndpoint, + body: Vec, + keep: bool, + ) -> Result { + self.sent_keeps.lock().unwrap().push(keep); self.bodies .lock() .unwrap() @@ -163,6 +173,12 @@ impl TurnHost for Fake { async fn keep_images(&self) -> Option { *self.keep.lock().unwrap() } + fn prewarm_keeps(&self, prefix: &[u64], payload: &Value) -> bool { + // Ce que l'hôte reçoit : les empreintes des messages envoyés et le + // corps de la requête. + assert_eq!(prefix.len(), payload["messages"].as_array().unwrap().len()); + *self.prewarm_match.lock().unwrap() + } async fn keep_image(&self, m: &Message) -> Option<(Message, i64)> { *self.keep_stores.lock().unwrap() += 1; if *self.keep_refuse.lock().unwrap() { @@ -1055,3 +1071,26 @@ async fn images_retirees_fenetre_pleine() { assert!(evs.iter().any(|e| matches!(e, StreamEvent::DropReasoning))); assert_eq!(text(&evs), "fini"); } + +/// PREWARM : la boucle demande à l'hôte si la requête prolonge le +/// préchauffage en vol et le lui transmet à l'envoi ; jamais vers une API +/// externe. +#[tokio::test] +async fn prechauffage_garde_par_la_requete_qui_le_prolonge() { + let host = Fake::new(vec![answer("ok")]); + *host.prewarm_match.lock().unwrap() = true; + let (r, _) = run(&host, &config(&[]), &Cancel::never()).await; + assert_eq!(r.err, None); + assert_eq!(*host.sent_keeps.lock().unwrap(), [true]); + let host = Fake::new(vec![answer("ok")]); + let (r, _) = run(&host, &config(&[]), &Cancel::never()).await; + assert_eq!(r.err, None); + assert_eq!(*host.sent_keeps.lock().unwrap(), [false]); + let host = Fake::new(vec![answer("ok")]); + *host.prewarm_match.lock().unwrap() = true; + let mut tc = config(&[]); + tc.ep.external = true; + let (r, _) = run(&host, &tc, &Cancel::never()).await; + assert_eq!(r.err, None); + assert_eq!(*host.sent_keeps.lock().unwrap(), [false]); +} diff --git a/rust/loki-rs/src/chatapi.rs b/rust/loki-rs/src/chatapi.rs index ae411d9..b3c6eec 100644 --- a/rust/loki-rs/src/chatapi.rs +++ b/rust/loki-rs/src/chatapi.rs @@ -66,6 +66,8 @@ pub fn is_internal(path: &str) -> bool { | "/internal/facade/v1/project/switch" | "/internal/facade/v1/project/drop" | "/internal/facade/v1/chat/tpl" + | "/internal/facade/v1/chat/prewarm" + | "/internal/facade/v1/chat/prewarm/cancel" ) } @@ -107,6 +109,27 @@ pub async fn serve_internal(chat: &Arc, req: Request) -> None => send_json(StatusCode::OK, &json!({ "ok": false })), } } + // PREWARM (lot 33) : déclencheur venu du Go (fin d'une tâche + // planifiée), annulation relayée (une autre requête est partie). + "/internal/facade/v1/chat/prewarm" => { + let v = read_json(req.into_body()).await.unwrap_or_default(); + let why = match str_of(&v, "why") { + crate::chat::prewarm::AFTER_TASK => crate::chat::prewarm::AFTER_TASK, + crate::chat::prewarm::SWITCH => crate::chat::prewarm::SWITCH, + _ => crate::chat::prewarm::TURN_END, + }; + chat.prewarm_kick(why).await; + send_json(StatusCode::OK, &json!({ "ok": true })) + } + "/internal/facade/v1/chat/prewarm/cancel" => { + let v = read_json(req.into_body()).await.unwrap_or_default(); + let id = v.get("id").and_then(Value::as_u64).unwrap_or(0); + let cancelled = chat.shared.prewarm.cancel_id(id); + send_json( + StatusCode::OK, + &json!({ "ok": true, "cancelled": cancelled }), + ) + } "/internal/facade/v1/project/switch" | "/internal/facade/v1/project/drop" => { let v = read_json(req.into_body()).await.unwrap_or_default(); let slug = str_of(&v, "slug"); @@ -238,8 +261,16 @@ pub async fn serve( // Corps illisible : champs vides, comme le `_ = Decode` du Go. "/api/conversations/switch" => { let v = read_json(body).await.unwrap_or_default(); + let before = chat.sessions.conv().id(); match chat.sessions.switch(str_of(&v, "id")).await { - Ok(()) => ok(), + Ok(()) => { + // PREWARM : le préchauffage de l'ancienne est annulé ; avec + // « full », la nouvelle est préparée si on y reste. + if chat.sessions.conv().id() != before { + chat.prewarm_after_switch().await; + } + ok() + } Err(e) => session_error(e), } } diff --git a/rust/loki-rs/tests/common/fakego.rs b/rust/loki-rs/tests/common/fakego.rs index 8eb6ffa..8e22d9a 100644 --- a/rust/loki-rs/tests/common/fakego.rs +++ b/rust/loki-rs/tests/common/fakego.rs @@ -137,6 +137,10 @@ pub struct FakeGo { /// /internal/v1/turn/keepimg : vision active ; rangements reçus. pub keep_vision: Mutex, pub keep_stores: Mutex>, + /// /internal/v1/prewarm/* : refus du contrôle (raison), appels reçus + /// (« check », « begin {corps} », « end {corps} »). + pub prewarm_refuse: Mutex>, + pub prewarm_log: Mutex>, } impl FakeGo { @@ -680,10 +684,15 @@ async fn handle(st: Arc, req: Request) -> Response let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); if let Some(conv) = v["main_conv"].as_str() { let side = if v["side"] == true { " side" } else { "" }; + let keep = if v["prewarm_keep"] == true { + " prewarm_keep" + } else { + "" + }; st.oai_log .lock() .unwrap() - .push(format!("begin main_conv={conv}{side}")); + .push(format!("begin main_conv={conv}{side}{keep}")); } else { st.oai_log.lock().unwrap().push(format!( "begin read={} side={} post_v1={}", @@ -695,6 +704,29 @@ async fn handle(st: Arc, req: Request) -> Response } resp(200, r#"{"ticket":7}"#) } + "/internal/v1/prewarm/check" => { + st.prewarm_log.lock().unwrap().push("check".into()); + let v = match st.prewarm_refuse.lock().unwrap().clone() { + Some(why) => serde_json::json!({ "ok": false, "why": why }), + None => serde_json::json!({ "ok": true, "slot": -1 }), + }; + resp(200, v.to_string()) + } + "/internal/v1/prewarm/begin" | "/internal/v1/prewarm/end" => { + let body = req.into_body().collect().await.unwrap().to_bytes(); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let what = if path.ends_with("begin") { + "begin" + } else { + "end" + }; + st.prewarm_log.lock().unwrap().push(format!("{what} {v}")); + if what == "begin" { + resp(200, r#"{"ok":true,"id":5}"#) + } else { + resp(204, "") + } + } "/internal/v1/turn/keepimg" => { let v = serde_json::json!({ "vision": *st.keep_vision.lock().unwrap(), diff --git a/rust/loki-rs/tests/lot6.rs b/rust/loki-rs/tests/lot6.rs index aef62fe..8921402 100644 --- a/rust/loki-rs/tests/lot6.rs +++ b/rust/loki-rs/tests/lot6.rs @@ -666,6 +666,7 @@ async fn tour_complet_par_le_go() { images: ImageCache::new(imgdir), patterns: RetryPatterns::builtin(), tpl: Arc::new(loki_rs::chat::tplprobe::TplProbe::new()), + prewarm: Default::default(), }); let cfg = std::collections::HashMap::new(); let host = GoTurnHost { @@ -2082,3 +2083,135 @@ async fn images_gardees_par_la_facade() { "image encore envoyée : {sent}" ); } + +/// Réponse non streamée d'un préchauffage (un seul jeton, jeté). +fn prewarm_reply() -> String { + serde_json::json!({ + "choices": [{"message": {"content": "x"}, "finish_reason": "length"}], + "usage": {"prompt_tokens": 60, "completion_tokens": 1}, + "timings": {"prompt_n": 18, "cache_n": 42} + }) + .to_string() +} + +/// Lot 33 : PREWARM tenu par la façade. Fin de tour → la requête du prochain +/// tour, sentinelle « . » et un seul jeton, inscrite au Go et rendue avec ce +/// que le moteur a repris. Le tour suivant, qui la prolonge, ne l'annule pas. +#[tokio::test] +async fn prechauffage_par_la_facade() { + let gate = Arc::new(tokio::sync::Semaphore::new(1)); + let (url, seen) = fake_engine_gated( + vec![answer("vu"), prewarm_reply(), answer("ok"), prewarm_reply()], + Some(gate.clone()), + ) + .await; + let (svc, go, _stop, _t) = service(Some(&url)).await; + go.set("config", "PREWARM", "on"); + svc.send(serde_json::json!({"message": "bonjour", "cid": "p1"})) + .await + .unwrap(); + until_done(&svc, 1).await; + let id = svc.sessions.conv().id(); + // Le préchauffage attend le moteur : il est en vol. + until(|| seen.lock().unwrap().len() == 2).await; + assert_eq!(go.prewarm_log.lock().unwrap()[..1], ["check".to_string()]); + let pre = seen.lock().unwrap()[1].1.clone(); + assert_eq!(pre["stream"], false); + assert_eq!(pre["max_tokens"], 1); + let msgs = pre["messages"].as_array().unwrap().clone(); + assert_eq!(msgs.last().unwrap()["content"], "."); + assert_eq!(msgs.last().unwrap()["role"], "user"); + let first = seen.lock().unwrap()[0].1.clone(); + assert_eq!(msgs[..2], first["messages"].as_array().unwrap()[..2]); + assert_eq!(msgs[2]["content"], "vu"); + assert_eq!(pre["tools"], first["tools"]); + let begin = go.prewarm_log.lock().unwrap()[1].clone(); + assert!( + begin.starts_with("begin ") + && begin.contains(&format!("\"conv\":\"{id}\"")) + && begin.contains("\"n\":3"), + "{begin}" + ); + // Pas de ticket « en vol » de la façade : le Go l'a inscrit lui-même. + let begins = || { + go.oai_log + .lock() + .unwrap() + .iter() + .filter(|l| l.starts_with("begin")) + .cloned() + .collect::>() + }; + assert_eq!(begins().len(), 1); + + // Le tour suivant prolonge le préfixe préchauffé : gardé. + svc.send(serde_json::json!({"message": "suite", "cid": "p2"})) + .await + .unwrap(); + until(|| begins().len() == 2).await; + assert_eq!(begins()[1], format!("begin main_conv={id} prewarm_keep")); + gate.add_permits(10); + until_done(&svc, 2).await; + until(|| { + go.prewarm_log + .lock() + .unwrap() + .iter() + .filter(|l| l.starts_with("end")) + .count() + == 2 + }) + .await; + let log = go.prewarm_log.lock().unwrap().clone(); + let first_end = log.iter().find(|l| l.starts_with("end")).unwrap(); + let v: serde_json::Value = serde_json::from_str(first_end.trim_start_matches("end ")).unwrap(); + assert_eq!(v, serde_json::json!({"id": 5, "total": 60, "cached": 42})); + let perfs = go.perf_records.lock().unwrap().clone(); + assert!( + perfs + .iter() + .any(|p| p["kind"] == "prewarm" && p["stats"]["prompt_tokens_total"] == 60) + ); +} + +/// Le Go refuse (moteur occupé) : rien ne part ; annulé par une autre +/// requête : la connexion est coupée et la place rendue. +#[tokio::test] +async fn prechauffage_refuse_ou_annule() { + let (url, seen) = fake_engine(vec![answer("vu")]).await; + let (svc, go, _stop, _t) = service(Some(&url)).await; + go.set("config", "PREWARM", "on"); + *go.prewarm_refuse.lock().unwrap() = Some("moteur occupé".into()); + svc.send(serde_json::json!({"message": "bonjour", "cid": "r1"})) + .await + .unwrap(); + until_done(&svc, 1).await; + until(|| !go.prewarm_log.lock().unwrap().is_empty()).await; + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(*go.prewarm_log.lock().unwrap(), ["check"]); + assert_eq!(seen.lock().unwrap().len(), 1); + + // Annulation relayée par le Go pendant que le moteur calcule. + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let (url, seen) = fake_engine_gated(vec![prewarm_reply()], Some(gate.clone())).await; + go.set("config", "PORT", &port_of(&url)); + *go.prewarm_refuse.lock().unwrap() = None; + go.prewarm_log.lock().unwrap().clear(); + svc.prewarm_kick(loki_rs::chat::prewarm::AFTER_TASK).await; + until(|| seen.lock().unwrap().len() == 1).await; + assert!(svc.shared.prewarm.live()); + assert!(svc.shared.prewarm.cancel_id(5)); + until(|| !svc.shared.prewarm.live()).await; + let log = go.prewarm_log.lock().unwrap().clone(); + assert_eq!( + log.last().unwrap(), + r#"end {"cached":null,"id":5,"total":0}"# + ); + assert!( + go.perf_records + .lock() + .unwrap() + .iter() + .all(|p| p["kind"] != "prewarm") + ); +} diff --git a/rust/loki-rs/tests/parity.rs b/rust/loki-rs/tests/parity.rs index 7f759ec..1e3be38 100644 --- a/rust/loki-rs/tests/parity.rs +++ b/rust/loki-rs/tests/parity.rs @@ -63,7 +63,12 @@ struct ParityHost { } impl TurnHost for ParityHost { - async fn send(&self, _ep: &ChatEndpoint, body: Vec) -> Result { + async fn send( + &self, + _ep: &ChatEndpoint, + body: Vec, + _keep: bool, + ) -> Result { self.requests .lock() .unwrap()