From bbea5fcd3f0bd8f8f0be2049f42070d9e32b5861 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 13:11:07 +0000 Subject: [PATCH] =?UTF-8?q?Lot=208=20:=20t=C3=A9l=C3=A9chargements=20de=20?= =?UTF-8?q?mod=C3=A8les=20tenus=20par=20la=20fa=C3=A7ade=20(LOKI=5FDOWNLOA?= =?UTF-8?q?DS=3Drust)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit La façade Rust mène les téléchargements de .gguf quand LOKI_DOWNLOADS=rust (expérimental ; le Go reste le défaut) : sonde, lancement, suivi, annulation, avec les mêmes contrats et la même mécanique que le Go (tranches rapatriées en entier, plages d'octets parallèles, reprise, place disque vérifiée sur le total, jeton réservé aux hôtes Hugging Face, refus traduits, .part vidés au démarrage). Côté Go : facade_downloads.go remet ces routes à la façade quand elles arrivent par le tunnel du relais, et le Go ne vide plus les .part qu'elle mène. Parité : journal des fonctions pures tiré du Go (parity_dl_test.go) et rejoué en Rust ; téléchargements de bout en bout contre un faux Hugging Face (tests/downloads.rs). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016tcGfJWefW5kUGoC2wL8cv --- README.md | 1 + ...t-migration-lot8-telechargements-design.md | 76 + internal/loki/facade_chat.go | 23 +- internal/loki/facade_chat_test.go | 75 + internal/loki/facade_downloads.go | 45 + internal/loki/parity_dl_test.go | 182 +++ internal/loki/web_server.go | 9 +- rust/Cargo.lock | 1 + rust/loki-rs/Cargo.toml | 1 + rust/loki-rs/src/api.rs | 3 + rust/loki-rs/src/chat/engine.rs | 5 +- rust/loki-rs/src/downloads.rs | 1306 +++++++++++++++++ rust/loki-rs/src/lib.rs | 1 + rust/loki-rs/src/main.rs | 45 + rust/loki-rs/src/proxy.rs | 15 +- rust/loki-rs/tests/common/fakego.rs | 41 +- rust/loki-rs/tests/downloads.rs | 692 +++++++++ rust/loki-rs/tests/lot2.rs | 1 + .../loki-rs/tests/parity_models/download.json | 314 ++++ rust/loki-rs/tests/proxy.rs | 1 + 20 files changed, 2820 insertions(+), 17 deletions(-) create mode 100644 docs/superpowers/specs/2026-10-15-rust-migration-lot8-telechargements-design.md create mode 100644 internal/loki/facade_downloads.go create mode 100644 internal/loki/parity_dl_test.go create mode 100644 rust/loki-rs/src/downloads.rs create mode 100644 rust/loki-rs/tests/downloads.rs create mode 100644 rust/loki-rs/tests/parity_models/download.json diff --git a/README.md b/README.md index 9b4fec9..98a1a9e 100644 --- a/README.md +++ b/README.md @@ -228,6 +228,7 @@ redémarrages (volume `/data`). Variables d'environnement du conteneur : | `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` | `rust` : la conversation est tenue par la façade Rust (expérimental, migration en cours), outils fichiers et shell compris. Le Go garde les autres outils, les pièces jointes, les fichiers et l'export. Avec `PROJ_SNAPSHOT`, `KEEP_TURN_IMAGES`, `NUDGE_IN_TOOL`, `SIDE_SLOT`, `SLOT_PERSIST`, `PREWARM` ou des motifs `retry_patterns`, le Go garde la conversation (décidé au démarrage) | Go | +| `LOKI_DOWNLOADS` | `rust` : les téléchargements de modèles (`/api/models/download*`) sont menés par la façade Rust (expérimental, migration en cours). Le Go garde les autres routes des modèles | Go | **Accès par nom de domaine.** Sans clé de pilotage, l'API n'accepte que les hôtes locaux (IP, `localhost`, nom sans point, `.local`/`.lan`…) et refuse diff --git a/docs/superpowers/specs/2026-10-15-rust-migration-lot8-telechargements-design.md b/docs/superpowers/specs/2026-10-15-rust-migration-lot8-telechargements-design.md new file mode 100644 index 0000000..4d74c70 --- /dev/null +++ b/docs/superpowers/specs/2026-10-15-rust-migration-lot8-telechargements-design.md @@ -0,0 +1,76 @@ +# Lot 8 — téléchargements de modèles tenus par la façade + +## Contexte + +Les téléchargements de `.gguf` (Hugging Face ou lien direct) vivaient dans le +Go : `backend_models.go`, `backend_models_shards.go`, `backend_model_dirs.go`, +`backend_hf.go`. Lot 8 : la façade les tient quand `LOKI_DOWNLOADS=rust` +(expérimental ; sans la variable rien ne change, tout part au Go). + +## Ce que sert la façade (`rust/loki-rs/src/downloads.rs`) + +| Route | Rôle | +|---|---| +| `POST /api/models/download/probe` | taille du modèle (toutes ses tranches), espace libre, « ça tient ? », ou « déjà là » avec la valeur de `MODEL=` | +| `POST /api/models/download` | lance en arrière-plan ; 409 si déjà en cours ; « déjà là » si toutes les tranches sont présentes | +| `GET /api/models/download/status` | état de chaque téléchargement de ce lancement | +| `POST /api/models/download/cancel` | annule ; ce que le téléchargement a écrit est effacé | + +Mêmes contrats (corps, codes, messages) et même mécanique que le Go : + +- lien de page Hugging Face (`/blob/`) converti en lien direct ; +- modèle en tranches rapatrié en entier (la première le désigne) ; +- tranches déjà là non retéléchargées ; +- place disque vérifiée sur le total, avec 256 Mio de marge ; jamais de refus + sur FUSE, overlayfs, NFS, SMB, Ceph ou GlusterFS, où l'espace annoncé + approxime ; +- jusqu'à `LOKI_DL_CONNS` connexions (8 par défaut, 16 au plus) par tranche + quand la source honore les plages, 16 Mio par connexion au moins ; +- connexion coupée reprise où elle en était (4 essais) ; +- progression sur l'ensemble, vitesse lissée ; +- jeton Hugging Face (réglage, sinon `HF_TOKEN`) envoyé aux seuls hôtes + Hugging Face ; +- refus de la source traduits en phrase qui dit quoi faire (dépôt verrouillé, + jeton absent ou insuffisant, fichier absent, quota…) ; +- `.part` orphelins vidés au démarrage — par la façade : le Go ne le fait + plus, il effacerait un téléchargement qu'elle mène. + +## Côté Go (`facade_downloads.go`) + +Le tunnel du relais et la redistribution E2E arrivent au Go sans passer par +la façade : `withFacadeChat` lui remet les quatre routes, authentifiées par le +jeton interne, comme celles du chat. Façade injoignable : 502 +« téléchargements indisponibles : … ». + +La façade transmet `LOKI_DOWNLOADS=rust` et son adresse au Go qu'elle lance. + +## Écarts assumés + +- Une source qui ignore les plages (200 à la sonde) n'est plus lue en entier + pendant la sonde. +- Erreur de transport (réseau, TLS) au format de reqwest. +- Hôte d'un lien collé en majuscules converti comme les autres. +- Espace libre inconnu hors Unix : rien n'est refusé (comme une mesure ratée + côté Go). +- `/status` trié par date de début (le Go rendait l'ordre d'une map). + +## Parité + +`internal/loki/parity_dl_test.go` tire du Go le journal des fonctions pures +(lien normalisé, famille de tranches, `Content-Range`, messages de refus, +hôtes et dépôts Hugging Face, dossier de destination), rangé dans +`rust/loki-rs/tests/parity_models/download.json` ; `tests/downloads.rs` le +rejoue (`LOKI_PARITY_UPDATE=1` pour régénérer). Une mutation +(`replacen` → `replace` sur `/blob/`) fait échouer le rejeu. + +`tests/downloads.rs` mène aussi le téléchargement de bout en bout contre un +faux Hugging Face : trois connexions sur des plages disjointes, source sans +plages, tranches avec une déjà présente, dépôt verrouillé refusé dès la sonde, +panne en cours de deuxième tranche, annulation, routes par la façade (clé de +pilotage, remise du Go), routes laissées au Go sans la variable. + +## Reste + +- Les autres routes des modèles (`/api/models`, suppression, dossiers, liste + Hugging Face, vision) restent au Go. +- Essai réel sur une installation avant d'en faire le défaut. diff --git a/internal/loki/facade_chat.go b/internal/loki/facade_chat.go index 1bd95de..225c809 100644 --- a/internal/loki/facade_chat.go +++ b/internal/loki/facade_chat.go @@ -141,23 +141,30 @@ var facadeChatRoutes = map[string]bool{ } // withFacadeChat remet la requête à la façade quand elle tient la -// conversation. Posé APRÈS requireWebAuth : seule une requête authentifiée -// ici part, avec le jeton interne qui en tient lieu là-bas. +// conversation — ou les téléchargements de modèles (facade_downloads.go). +// Posé APRÈS requireWebAuth : seule une requête authentifiée ici part, avec le +// jeton interne qui en tient lieu là-bas. func withFacadeChat(path string, h http.HandlerFunc) http.HandlerFunc { - if !facadeChatRoutes[path] { + held, unavailable := facadeChat, "conversation indisponible : " + switch { + case facadeChatRoutes[path]: + case facadeDownloadRoutes[path]: + held, unavailable = facadeDownloads, "téléchargements indisponibles : " + default: return h } return func(w http.ResponseWriter, r *http.Request) { - if !facadeChat() { + if !held() { h(w, r) return } - facadeRelay(w, r) + facadeRelayAs(w, r, unavailable) } } -// facadeRelay : relais transparent (flux compris) vers la façade. -func facadeRelay(w http.ResponseWriter, r *http.Request) { +// facadeRelayAs : relais transparent (flux compris) vers la façade ; +// unavailable préfixe l'erreur si elle ne répond pas. +func facadeRelayAs(w http.ResponseWriter, r *http.Request, unavailable string) { rp := &httputil.ReverseProxy{ Rewrite: func(pr *httputil.ProxyRequest) { pr.SetURL(&url.URL{Scheme: "http", Host: facadeAddr()}) @@ -165,7 +172,7 @@ func facadeRelay(w http.ResponseWriter, r *http.Request) { }, FlushInterval: -1, // SSE : chaque trame part aussitôt ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) { - sendJSON(w, http.StatusBadGateway, map[string]any{"ok": false, "error": "conversation indisponible : " + err.Error()}) + sendJSON(w, http.StatusBadGateway, map[string]any{"ok": false, "error": unavailable + err.Error()}) }, } rp.ServeHTTP(w, r) diff --git a/internal/loki/facade_chat_test.go b/internal/loki/facade_chat_test.go index 70282f1..9b8ee5d 100644 --- a/internal/loki/facade_chat_test.go +++ b/internal/loki/facade_chat_test.go @@ -5,6 +5,8 @@ import ( "io" "net/http" "net/http/httptest" + "os" + "path/filepath" "strings" "sync" "testing" @@ -303,3 +305,76 @@ func TestFacadeReglageQuiGardeLaConversation(t *testing.T) { t.Fatalf("sans la variable : %s", w.Body.String()) } } + +// LOKI_DOWNLOADS=rust : les routes des téléchargements de modèles partent à la +// façade (avec ou sans LOKI_CHAT), les autres routes des modèles restent ici ; +// sans la variable, rien ne change. Façade muette : erreur qui le dit. +func TestFacadeRelaieLesTelechargements(t *testing.T) { + testHome(t) + f := demarrerFacade(t) + t.Setenv(facadeChatEnv, "") + t.Setenv(facadeDownloadsEnv, "rust") + local := func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "go") } + + for _, p := range []string{"/api/models/download", "/api/models/download/probe", + "/api/models/download/status", "/api/models/download/cancel"} { + w := httptest.NewRecorder() + withFacadeChat(p, local)(w, httptest.NewRequest("POST", p, strings.NewReader(`{"url":"u"}`))) + if w.Body.String() != f.reply { + t.Fatalf("%s : réponse %q", p, w.Body.String()) + } + } + if got := f.recu(); len(got) != 4 || got[0] != `POST /api/models/download jeton {"url":"u"}` { + t.Fatalf("reçu par la façade : %v", got) + } + // Le chat, lui, reste ici (LOKI_CHAT absent), et les autres routes des + // modèles aussi. + for _, p := range []string{"/api/chat/send", "/api/models", "/api/models/delete"} { + w := httptest.NewRecorder() + withFacadeChat(p, local)(w, httptest.NewRequest("POST", p, nil)) + if w.Body.String() != "go" { + t.Fatalf("%s relayé : %q", p, w.Body.String()) + } + } + // Façade injoignable : 502 qui nomme ce qui manque. + f.srv.Close() + w := httptest.NewRecorder() + withFacadeChat("/api/models/download/status", local)(w, httptest.NewRequest("GET", "/api/models/download/status", nil)) + if w.Code != http.StatusBadGateway || !strings.Contains(w.Body.String(), "téléchargements indisponibles : ") { + t.Fatalf("façade muette : %d %s", w.Code, w.Body.String()) + } + // Sans LOKI_DOWNLOADS=rust : tout reste ici. + t.Setenv(facadeDownloadsEnv, "") + w = httptest.NewRecorder() + withFacadeChat("/api/models/download", local)(w, httptest.NewRequest("POST", "/api/models/download", nil)) + if w.Body.String() != "go" { + t.Fatalf("relayé sans la variable : %q", w.Body.String()) + } +} + +// Les .part orphelins ne sont plus vidés par ce Go quand la façade tient les +// téléchargements : elle les vide à son démarrage, et un Go relancé seul +// effacerait celui qu'elle mène. +func TestFacadeTelechargementsGardeLesPart(t *testing.T) { + home := testHome(t) + dir := filepath.Join(home, "models") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + part := filepath.Join(dir, "m.gguf.part") + if err := os.WriteFile(part, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + t.Setenv(internalTokenEnv, jetonTest) + t.Setenv(facadeAddrEnv, "127.0.0.1:1") + t.Setenv(facadeDownloadsEnv, "rust") + cleanStartupPartFiles() + if _, err := os.Stat(part); err != nil { + t.Fatalf(".part effacé alors que la façade tient les téléchargements : %v", err) + } + t.Setenv(facadeDownloadsEnv, "") + cleanStartupPartFiles() + if _, err := os.Stat(part); !os.IsNotExist(err) { + t.Fatalf(".part orphelin gardé : %v", err) + } +} diff --git a/internal/loki/facade_downloads.go b/internal/loki/facade_downloads.go new file mode 100644 index 0000000..c35e074 --- /dev/null +++ b/internal/loki/facade_downloads.go @@ -0,0 +1,45 @@ +package loki + +import ( + "os" + "strings" +) + +// facade_downloads.go — téléchargements de modèles tenus par la façade Rust +// (LOKI_DOWNLOADS=rust, lot « modèles » de la migration ; expérimental, le Go +// reste le défaut). +// +// La façade sert elle-même /api/models/download* sur le port public. Le tunnel +// du relais et sa redistribution E2E arrivent pourtant ICI : ces routes lui +// sont remises, sans quoi un téléchargement lancé depuis le téléphone vivrait +// dans ce process, que le poste — servi par la façade — ne verrait pas. +// +// Les .part orphelins sont vidés par la façade à son démarrage : les vider ici +// (un Go relancé seul après un plantage) effacerait sous ses pieds un +// téléchargement qu'elle mène. + +const facadeDownloadsEnv = "LOKI_DOWNLOADS" + +// facadeDownloads : la façade tient-elle les téléchargements de modèles ? Il +// faut aussi son adresse et le jeton interne, sans lesquels rien ne peut lui +// être remis. +func facadeDownloads() bool { + return strings.EqualFold(strings.TrimSpace(os.Getenv(facadeDownloadsEnv)), "rust") && + facadeAddr() != "" && internalToken() != "" +} + +// cleanStartupPartFiles vide les .part orphelins au démarrage — sauf quand la +// façade tient les téléchargements : elle le fait à SON démarrage. +func cleanStartupPartFiles() { + if !facadeDownloads() { + cleanStalePartFiles() + } +} + +// facadeDownloadRoutes : les routes que la façade sert alors (downloads.rs). +var facadeDownloadRoutes = map[string]bool{ + "/api/models/download": true, + "/api/models/download/probe": true, + "/api/models/download/status": true, + "/api/models/download/cancel": true, +} diff --git a/internal/loki/parity_dl_test.go b/internal/loki/parity_dl_test.go new file mode 100644 index 0000000..784c577 --- /dev/null +++ b/internal/loki/parity_dl_test.go @@ -0,0 +1,182 @@ +package loki + +import ( + "net/http" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +// parity_dl_test.go — journal de référence des fonctions pures des +// téléchargements de modèles (lot « modèles » de la migration Rust) : lien +// collé normalisé, famille de tranches, messages d'erreur de la source, +// dossier de destination. Rangé dans +// rust/loki-rs/tests/parity_models/download.json ; le test Rust +// (tests/downloads.rs) doit retrouver chaque sortie. LOKI_HOME réécrit en +// {{HOME}}. +// +// LOKI_PARITY_UPDATE=1 go test ./internal/loki -run TestPariteTelechargements + +const parityDLFile = "../../rust/loki-rs/tests/parity_models/download.json" + +type parityDLNorm struct { + In string `json:"in"` + URL string `json:"url,omitempty"` + Name string `json:"name,omitempty"` + Err string `json:"err,omitempty"` + URLs []string `json:"urls,omitempty"` + // Names : noms de la famille de tranches (shardURLSet). + Names []string `json:"names,omitempty"` +} + +type parityDLSource struct { + Status int `json:"status"` + Code string `json:"code"` + URL string `json:"url"` + TokenSet bool `json:"token_set"` + Err string `json:"err"` +} + +type parityDLDest struct { + Name string `json:"name"` + Dir string `json:"dir"` + Path string `json:"path,omitempty"` + Err string `json:"err,omitempty"` +} + +type parityDLFixture struct { + Normalize []parityDLNorm `json:"normalize"` + ContentRange map[string]int64 `json:"content_range"` + Source []parityDLSource `json:"source"` + HFHost map[string]bool `json:"hf_host"` + HFRepo map[string]string `json:"hf_repo"` + Dest []parityDLDest `json:"dest"` +} + +func TestPariteTelechargements(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("journal tiré sous Unix") + } + home := testHome(t) + if r, err := filepath.EvalSymlinks(home); err == nil { + home = r + } + t.Setenv("LOKI_HOME", home) + t.Setenv("LOKI_MODEL_DIRS", "") + t.Setenv("HF_TOKEN", "") + norm := func(s string) string { return strings.ReplaceAll(s, home, "{{HOME}}") } + + var fx parityDLFixture + for _, in := range []string{ + "", " ", "ftp://h/a.gguf", "pas un lien", "https://exemple.org/dossier/", + "https://huggingface.co/u/r/blob/main/m.gguf", + "https://huggingface.co/u/r/blob/main/blob/m.gguf", + "https://huggingface.co/u/r/resolve/main/m-Q4_K_M.gguf?download=true", + " https://huggingface.co/u/r/blob/main/m.bin ", + "http://exemple.org/modeles/M.GGUF", + "https://huggingface.co/u/r/blob/main/Q4/m-Q4-00002-of-00003.gguf", + "https://huggingface.co/u/r/resolve/main/m-00001-of-00012.gguf", + "https://exemple.org/m-1-of-2.gguf", + "https://exemple.org/m-02-of-01.gguf", + } { + c := parityDLNorm{In: in} + u, name, err := normalizeHFURL(in) + if err != nil { + c.Err = err.Error() + } else { + c.URL, c.Name = u, name + c.URLs, c.Names = shardURLSet(u, name) + } + fx.Normalize = append(fx.Normalize, c) + } + + fx.ContentRange = map[string]int64{} + for _, v := range []string{"bytes 0-0/12345", "bytes 0-0/*", "", "bytes 0-0/0", "bytes 0-0/ 99 "} { + fx.ContentRange[v] = contentRangeTotal(v) + } + + hf := "https://huggingface.co/auteur/depot/resolve/main/m.gguf" + for _, c := range []parityDLSource{ + {Status: 401, Code: "GatedRepo", URL: hf}, + {Status: 401, Code: "GatedRepo", URL: hf, TokenSet: true}, + {Status: 403, Code: "RepoNotFound", URL: hf}, + {Status: 401, Code: "GatedRepo", URL: "https://exemple.org/m.gguf"}, + {Status: 403, Code: "", URL: hf, TokenSet: true}, + {Status: 403, Code: "", URL: "https://exemple.org/m.gguf"}, + {Status: 404, Code: "EntryNotFound", URL: hf}, + {Status: 404, Code: "EntryNotFound", URL: "https://exemple.org/"}, + {Status: 404, Code: "RevisionNotFound", URL: hf}, + {Status: 404, Code: "", URL: hf}, + {Status: 416, URL: hf}, {Status: 429, URL: hf}, {Status: 502, URL: hf}, {Status: 418, URL: hf}, + } { + tok := "" + if c.TokenSet { + tok = "hf_x" + } + t.Setenv("HF_TOKEN", tok) + resp := &http.Response{StatusCode: c.Status, Header: http.Header{}} + if c.Code != "" { + resp.Header.Set("X-Error-Code", c.Code) + } + c.Err = dlSourceError(resp, c.URL).Error() + fx.Source = append(fx.Source, c) + } + t.Setenv("HF_TOKEN", "") + + fx.HFHost = map[string]bool{} + for _, h := range []string{"huggingface.co", "HUGGINGFACE.CO:443", "cdn-lfs.huggingface.co", "hf.co", "x.hf.co", + "evilhuggingface.co", "huggingface.co.evil.org", "exemple.org"} { + fx.HFHost[h] = isHFHost(h) + } + fx.HFRepo = map[string]string{} + for _, u := range []string{hf, "https://huggingface.co/a/b/blob/main/x.gguf", "https://huggingface.co/a/b/tree/main", + "https://huggingface.co/a/b", "https://exemple.org/a/b/resolve/main/x.gguf", "https://huggingface.co/-a/b/resolve/main/x.gguf"} { + fx.HFRepo[u] = hfRepoFromURL(u) + } + + extra := filepath.Join(home, "disque") + if err := os.MkdirAll(extra, 0o755); err != nil { + t.Fatal(err) + } + if err := saveExtraModelDirs([]string{extra}); err != nil { + t.Fatal(err) + } + for _, c := range []parityDLDest{ + {Name: "a.gguf"}, {Name: " ../../x.gguf "}, {Name: "x.bin"}, {Name: ""}, {Name: "/"}, + {Name: "a.gguf", Dir: "/etc"}, {Name: "a.gguf", Dir: "{{HOME}}/disque"}, {Name: "a.gguf", Dir: "{{HOME}}/disque/"}, + {Name: "a.gguf", Dir: "{{HOME}}/models"}, + } { + p, err := downloadDestPath(c.Name, strings.ReplaceAll(c.Dir, "{{HOME}}", home)) + if err != nil { + c.Err = norm(err.Error()) + } else { + c.Path = norm(p) + } + fx.Dest = append(fx.Dest, c) + } + + got, err := parityJSON(fx) + if err != nil { + t.Fatal(err) + } + if os.Getenv("LOKI_PARITY_UPDATE") == "1" { + if err := os.MkdirAll(filepath.Dir(parityDLFile), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(parityDLFile, got, 0o644); err != nil { + t.Fatal(err) + } + return + } + want, err := os.ReadFile(parityDLFile) + if err != nil { + t.Fatalf("journal de référence absent (%v) : LOKI_PARITY_UPDATE=1 pour le créer", err) + } + if string(got) != string(want) { + t.Errorf("les téléchargements ne produisent plus le journal rangé dans %s.\n"+ + "Si c'est voulu : LOKI_PARITY_UPDATE=1 go test ./internal/loki -run TestPariteTelechargements,\n"+ + "puis cargo test --test downloads dit si le port Rust suit.\n--- obtenu ---\n%s", parityDLFile, got) + } +} diff --git a/internal/loki/web_server.go b/internal/loki/web_server.go index 4609ff9..99d3ffc 100644 --- a/internal/loki/web_server.go +++ b/internal/loki/web_server.go @@ -130,8 +130,9 @@ func newWebMux() *http.ServeMux { // secondes pour un serveur lancé via npx) est payé par le premier message. MCPPrewarm() // Un téléchargement de modèle coupé net (crash, restart du service) laisse un - // .part orphelin non reprenable : on nettoie au démarrage. - cleanStalePartFiles() + // .part orphelin non reprenable : on nettoie au démarrage (sauf quand la + // façade tient les téléchargements, facade_downloads.go). + cleanStartupPartFiles() // Les fichiers des versions précédentes vivaient dans un pot commun : on les // range dans le dossier de leur discussion (une fois par process). migrateConvFiles() @@ -218,7 +219,9 @@ func newWebMux() *http.ServeMux { mountOAI(mux) // api enregistre une route /api/* protégée par la clé de pilotage (web_auth.go). // withFacadeChat : routes du chat remises à la façade quand elle tient la - // conversation (LOKI_CHAT=rust, facade_chat.go) ; sans effet sinon. + // conversation (LOKI_CHAT=rust, facade_chat.go), celles des téléchargements + // de modèles quand elle les tient (LOKI_DOWNLOADS=rust, + // facade_downloads.go) ; sans effet sinon. api := func(path string, h http.HandlerFunc) { mux.HandleFunc(path, requireWebAuth(withFacadeChat(path, h))) } api("/api/ping", handlePing) api("/api/status", handleStatus) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 9df7f87..c003d08 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -686,6 +686,7 @@ dependencies = [ "tempfile", "tokio", "tokio-tungstenite", + "url", ] [[package]] diff --git a/rust/loki-rs/Cargo.toml b/rust/loki-rs/Cargo.toml index fcf2da2..191251e 100644 --- a/rust/loki-rs/Cargo.toml +++ b/rust/loki-rs/Cargo.toml @@ -13,6 +13,7 @@ http-body-util = "0.1.5" hyper = { version = "1.12.0", features = ["http1", "server", "client"] } hyper-util = { version = "0.1.21", features = ["tokio", "server", "client-legacy", "http1"] } regex = "1.13.1" +url = "2.5.8" reqwest = { version = "0.13.5", default-features = false, features = ["rustls-no-provider", "stream", "http2"] } rustls = { version = "0.23.46", default-features = false, features = ["ring", "std", "tls12"], optional = true } serde = { version = "1.0.229", features = ["derive"] } diff --git a/rust/loki-rs/src/api.rs b/rust/loki-rs/src/api.rs index 020a999..5caf66c 100644 --- a/rust/loki-rs/src/api.rs +++ b/rust/loki-rs/src/api.rs @@ -26,6 +26,9 @@ pub struct Ctx { pub engine: engine::Paths, /// Conversation tenue par la façade (`LOKI_CHAT=rust`) ; `None` : au Go. pub chat: Option>, + /// Téléchargements de modèles tenus par la façade (`LOKI_DOWNLOADS=rust`) ; + /// `None` : au Go. + pub downloads: Option>, /// Jeton de l'API interne : une requête qui le porte vient du Go. pub internal_token: String, } diff --git a/rust/loki-rs/src/chat/engine.rs b/rust/loki-rs/src/chat/engine.rs index fec61b5..205a9a2 100644 --- a/rust/loki-rs/src/chat/engine.rs +++ b/rust/loki-rs/src/chat/engine.rs @@ -208,8 +208,9 @@ impl EngineHttp { } } -/// Fournisseur cryptographique de rustls (ring), posé une fois. -fn install_crypto() { +/// Fournisseur cryptographique de rustls (ring), posé une fois (tout client +/// reqwest de la façade le demande avant d'être construit). +pub fn install_crypto() { #[cfg(feature = "tls")] let _ = rustls::crypto::ring::default_provider().install_default(); } diff --git a/rust/loki-rs/src/downloads.rs b/rust/loki-rs/src/downloads.rs new file mode 100644 index 0000000..08b84a5 --- /dev/null +++ b/rust/loki-rs/src/downloads.rs @@ -0,0 +1,1306 @@ +//! Téléchargements de modèles depuis Hugging Face (ou un lien direct) — port +//! des téléchargements de `backend_models.go`, `backend_models_shards.go`, +//! `backend_model_dirs.go` et `backend_hf.go` (lot 8 de la migration). +//! +//! Tenus par la façade quand `LOKI_DOWNLOADS=rust` (expérimental ; le Go +//! par défaut) : `POST /api/models/download`, `/probe`, `/cancel`, +//! `GET /api/models/download/status`. Le Go remet alors à la façade ces +//! routes quand elles lui arrivent par le tunnel du relais, et ne vide plus +//! les `.part` au démarrage (c'est la façade qui le fait). +//! +//! Comme le Go : un modèle découpé en tranches se rapatrie EN ENTIER (la +//! première désigne le modèle), les tranches déjà là ne se retéléchargent +//! pas, la place disque se vérifie sur le total, chaque tranche se répartit +//! sur plusieurs connexions quand la source honore les plages d'octets, une +//! connexion coupée reprend où elle en était, et une annulation ou un échec +//! efface ce que CE téléchargement avait écrit. +//! +//! Écarts assumés : une source qui ignore les plages (200 à la sonde) n'est +//! plus lue en entier pour rien pendant la sonde ; message d'une erreur de +//! transport (réseau, TLS) au format de reqwest ; hôte d'un lien collé en +//! majuscules converti comme les autres (`/blob/` → `/resolve/`) ; espace +//! libre inconnu hors Unix (rien n'est alors refusé, comme pour une mesure +//! ratée côté Go) ; `/status` trié par date de début (le Go rendait l'ordre +//! d'une map). + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use futures_util::StreamExt; +use http_body_util::{BodyExt, Limited}; +use hyper::body::Incoming; +use hyper::{Request, Response, StatusCode}; +use serde::Serialize; +use serde_json::{Value, json}; +use tokio::sync::watch; + +use crate::api::Ctx; +use crate::auth::json as send_json; +use crate::gguf; +use crate::paths::{base_name, go_abs, go_clean, go_join}; +use crate::proxy::Body; +use crate::store::{BK_STATE, Store}; + +/// Connexions parallèles par tranche (`LOKI_DL_CONNS`, 16 au plus). +pub fn dl_conns() -> usize { + std::env::var("LOKI_DL_CONNS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|n| *n > 0) + .unwrap_or(8) + .min(16) +} + +/// Plus petite part qui mérite sa propre connexion (16 Mio). +const DL_MIN_CHUNK: i64 = 16 << 20; +/// Marge laissée libre après le téléchargement. +const DL_SPACE_MARGIN: i64 = 256 << 20; +/// Essais d'une part avant d'abandonner. +const DL_ATTEMPTS: u64 = 4; +/// Écrit par blocs de 1 Mio. +const DL_WRITE_BUF: usize = 1 << 20; + +/// Routes servies ici quand la façade tient les téléchargements. +pub fn handles(path: &str) -> bool { + matches!( + path, + "/api/models/download" + | "/api/models/download/probe" + | "/api/models/download/status" + | "/api/models/download/cancel" + ) +} + +// --------------------------------------------------------------------------- +// Fonctions pures (journal de référence : tests/parity_models/download.json) +// --------------------------------------------------------------------------- + +/// `path.Base` sur un chemin d'URL (séparateur `/` seul). +fn url_path_base(p: &str) -> String { + if p.is_empty() { + return ".".into(); + } + let t = p.trim_end_matches('/'); + if t.is_empty() { + return "/".into(); + } + t.rsplit('/').next().unwrap_or(t).to_string() +} + +/// Décodage `%XX` d'un segment d'URL (le `u.Path` du Go est décodé). +fn percent_decode(s: &str) -> String { + let b = s.as_bytes(); + let mut out = Vec::with_capacity(b.len()); + let mut i = 0; + while i < b.len() { + if b[i] == b'%' + && i + 2 < b.len() + && let (Some(h), Some(l)) = (hex(b[i + 1]), hex(b[i + 2])) + { + out.push(h << 4 | l); + i += 3; + continue; + } + out.push(b[i]); + i += 1; + } + String::from_utf8_lossy(&out).into_owned() +} + +fn hex(c: u8) -> Option { + match c { + b'0'..=b'9' => Some(c - b'0'), + b'a'..=b'f' => Some(c - b'a' + 10), + b'A'..=b'F' => Some(c - b'A' + 10), + _ => None, + } +} + +/// `normalizeHFURL` : lien de page Hugging Face (`/blob/`) → lien direct +/// (`/resolve/`) ; rend l'URL à télécharger et le nom du fichier. +pub fn normalize_hf_url(raw: &str) -> Result<(String, String), String> { + let raw = raw.trim(); + if raw.is_empty() { + return Err("lien vide".into()); + } + let mut u = match reqwest::Url::parse(raw) { + Ok(u) => u, + // Pas de schéma : le Go l'accepte et refuse ensuite le schéma vide. + Err(url::ParseError::RelativeUrlWithoutBase) => { + return Err("lien invalide (http/https attendu)".into()); + } + Err(e) => return Err(format!("lien invalide: {e}")), + }; + if u.scheme() != "http" && u.scheme() != "https" { + return Err("lien invalide (http/https attendu)".into()); + } + if u.host_str().is_some_and(|h| h.contains("huggingface.co")) { + let p = u.path().replacen("/blob/", "/resolve/", 1); + u.set_path(&p); + } + let name = percent_decode(&url_path_base(u.path())); + if name.is_empty() || name == "/" || name == "." { + return Err("impossible de déduire le nom du fichier depuis le lien".into()); + } + if !name.to_lowercase().ends_with(".gguf") { + return Err("le lien doit pointer vers un fichier .gguf".into()); + } + Ok((u.to_string(), name)) +} + +/// `shardURLSet` : l'URL d'UNE tranche → les URLs de toute la famille, dans +/// l'ordre, avec leurs noms. +pub fn shard_url_set(dl_url: &str, name: &str) -> (Vec, Vec) { + let fam = gguf::shard_family(name); + if fam.len() == 1 { + return (vec![dl_url.to_string()], fam); + } + let base = base_name(name); + let Some(i) = dl_url.rfind(&base) else { + return (vec![dl_url.to_string()], vec![base]); + }; + let urls = fam + .iter() + .map(|n| format!("{}{n}{}", &dl_url[..i], &dl_url[i + base.len()..])) + .collect(); + (urls, fam) +} + +/// `contentRangeTotal` : la taille totale d'un « bytes 0-0/12345 ». +pub fn content_range_total(v: &str) -> i64 { + match v.rfind('/') { + Some(i) => v[i + 1..] + .trim() + .parse::() + .ok() + .filter(|n| *n > 0) + .unwrap_or(0), + None => 0, + } +} + +/// `isHFHost` : l'hôte appartient à Hugging Face (seul à recevoir le jeton). +pub fn is_hf_host(host: &str) -> bool { + let h = host.to_lowercase(); + let h = h.split(':').next().unwrap_or(""); + h == "huggingface.co" || h.ends_with(".huggingface.co") || h == "hf.co" || h.ends_with(".hf.co") +} + +fn hf_repo_ok(repo: &str) -> bool { + let part = |s: &str| { + let mut c = s.chars(); + c.next().is_some_and(|f| f.is_ascii_alphanumeric()) + && c.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '_' | '-')) + }; + match repo.split_once('/') { + Some((a, b)) => part(a) && part(b), + None => false, + } +} + +/// `hfRepoFromURL` : « auteur/dépôt » d'un lien de téléchargement, "" ailleurs. +pub fn hf_repo_from_url(raw: &str) -> String { + let Ok(u) = reqwest::Url::parse(raw) else { + return String::new(); + }; + if !u.host_str().unwrap_or("").contains("huggingface.co") { + return String::new(); + } + let segs: Vec<&str> = u.path().trim_matches('/').split('/').collect(); + if segs.len() < 4 || (segs[2] != "resolve" && segs[2] != "blob") { + return String::new(); + } + let repo = format!("{}/{}", segs[0], segs[1]); + if hf_repo_ok(&repo) { + repo + } else { + String::new() + } +} + +/// `dlFileLabel` : « nom.gguf » pour un message, "" sans nom lisible. +fn dl_file_label(dl_url: &str) -> String { + let Ok(u) = reqwest::Url::parse(dl_url) else { + return String::new(); + }; + let name = percent_decode(&url_path_base(u.path())); + if name.is_empty() || name == "/" || name == "." { + return String::new(); + } + format!("« {name} »") +} + +fn dl_access_reason(code: &str, repo: &str) -> String { + match (code, repo.is_empty()) { + ("GatedRepo", true) => { + "dépôt à accès restreint : ses conditions doivent être acceptées sur huggingface.co" + .into() + } + ("GatedRepo", false) => { + format!("dépôt à accès restreint : accepte ses conditions sur huggingface.co/{repo}") + } + ("RepoNotFound", true) => "dépôt privé ou inexistant".into(), + ("RepoNotFound", false) => format!("dépôt privé ou inexistant : {repo}"), + (_, false) => format!("accès refusé par Hugging Face sur {repo}"), + (_, true) => "accès refusé par la source".into(), + } +} + +fn dl_token_hint(token_set: bool) -> &'static str { + if token_set { + " — le jeton Hugging Face enregistré n'y donne pas accès (expiré, ou conditions non acceptées avec ce compte)" + } else { + " — puis renseigne le jeton Hugging Face (réglage « Jeton Hugging Face », ou variable d'environnement HF_TOKEN)" + } +} + +/// `dlSourceError` : un refus HTTP en phrase qui dit quoi faire. +pub fn dl_source_error(status: u16, error_code: &str, dl_url: &str, token_set: bool) -> String { + match status { + 401 | 403 => format!( + "{}{}", + dl_access_reason(error_code, &hf_repo_from_url(dl_url)), + dl_token_hint(token_set) + ), + 404 => match error_code { + "EntryNotFound" => format!( + "fichier {} absent du dépôt (HTTP 404) — dépôt réorganisé depuis la copie du lien, ou fragment jamais publié (table PLE livrée à part ?)", + dl_file_label(dl_url) + ), + "RevisionNotFound" => "révision introuvable dans le dépôt (HTTP 404)".into(), + _ => "lien introuvable (HTTP 404)".into(), + }, + 416 => "la source refuse la plage d'octets demandée (HTTP 416) — fichier modifié pendant le transfert ?".into(), + 429 => "trop de requêtes vers la source (HTTP 429) — réessaie dans quelques minutes".into(), + s if s >= 500 => format!("la source est en panne (HTTP {s}) — réessaie plus tard"), + s => format!("HTTP {s} depuis la source"), + } +} + +/// `humanBytes`. +pub fn human_bytes(n: i64) -> String { + if n >= 1 << 30 { + format!("{:.1} Go", n as f64 / (1u64 << 30) as f64) + } else if n >= 1 << 20 { + format!("{:.0} Mo", n as f64 / (1u64 << 20) as f64) + } else { + format!("{n} o") + } +} + +/// `filepath.Base` (séparateur de la plateforme). +fn filepath_base(p: &str) -> String { + if p.is_empty() { + return ".".into(); + } + let seps: &[char] = if cfg!(windows) { &['/', '\\'] } else { &['/'] }; + let t = p.trim_end_matches(seps); + if t.is_empty() { + return std::path::MAIN_SEPARATOR.to_string(); + } + t.rsplit(seps).next().unwrap_or(t).to_string() +} + +/// Les dossiers de modèles du moment (`modelsDir` + déclarés + variable). +#[derive(Debug, Clone)] +pub struct ModelDirs { + pub home: PathBuf, + pub stored: Vec, + pub env: Option, +} + +impl ModelDirs { + pub fn models_dir(&self) -> String { + gguf::models_dir(&self.home) + } + + pub fn all(&self) -> Vec { + gguf::model_dirs(&self.home, &self.stored, self.env.as_deref()) + } + + /// `resolveDownloadDir` : vide = models/, sinon un des dossiers déclarés. + pub fn resolve_download_dir(&self, dir: &str) -> Result { + let dir = dir.trim(); + if dir.is_empty() { + return Ok(self.models_dir()); + } + let abs = go_abs(dir).ok_or_else(|| "chemin invalide".to_string())?; + self.all() + .into_iter() + .find(|d| gguf::norm_dir(d) == gguf::norm_dir(&abs)) + .ok_or_else(|| { + format!("dossier non autorisé : {abs} — ajoute-le dans « Dossiers de modèles »") + }) + } + + /// `downloadDestPath`. + pub fn download_dest_path(&self, name: &str, dir: &str) -> Result { + let base = filepath_base(name.trim()); + if base.is_empty() || base == "." || base == std::path::MAIN_SEPARATOR_STR { + return Err("nom de modèle invalide".into()); + } + if !base.to_lowercase().ends_with(".gguf") { + return Err("seuls les fichiers .gguf sont acceptés".into()); + } + let d = self.resolve_download_dir(dir)?; + Ok(go_join(&[&d, &base])) + } + + /// `modelPickerValue` : le nom seul dans models/ (preset portable), le + /// chemin complet ailleurs. + pub fn model_picker_value(&self, dir: &str, base: &str) -> String { + if gguf::norm_dir(dir) == gguf::norm_dir(&self.models_dir()) { + base.to_string() + } else { + go_join(&[dir, base]) + } + } + + /// `modelPresence` : déjà sur le disque, famille complète ? Rend la + /// valeur de MODEL= et le chemin. + pub fn model_presence(&self, name: &str, dir: &str) -> Option<(String, String)> { + let base = base_name(name.trim()); + if base.is_empty() || !base.to_lowercase().ends_with(".gguf") { + return None; + } + let dirs = match self.resolve_download_dir(dir) { + Ok(d) if !dir.trim().is_empty() => vec![d], + _ => self.all(), + }; + for d in dirs { + let p = go_join(&[&d, &base]); + if !present(&p) { + continue; + } + // Famille incomplète : le téléchargement doit pouvoir la compléter. + if gguf::shard_family(&base) + .iter() + .any(|n| !present(&go_join(&[&d, n]))) + { + continue; + } + return Some((self.model_picker_value(&d, &base), p)); + } + None + } +} + +/// `os.Stat` sans erreur et pas un dossier. +fn present(p: &str) -> bool { + std::fs::metadata(p).is_ok_and(|m| !m.is_dir()) +} + +/// `modelFamilySize` : toutes les tranches présentes. +pub fn model_family_size(p: &str) -> i64 { + let path = Path::new(p); + let dir = path.parent().unwrap_or(Path::new(".")); + gguf::shard_family(&base_name(p)) + .iter() + .filter_map(|n| std::fs::metadata(dir.join(n)).ok()) + .filter(|m| !m.is_dir()) + .map(|m| m.len() as i64) + .sum() +} + +// --------------------------------------------------------------------------- +// Espace disque +// --------------------------------------------------------------------------- + +#[cfg(unix)] +fn statfs(dir: &str) -> Option { + let c = std::ffi::CString::new(dir).ok()?; + let mut st: libc::statfs = unsafe { std::mem::zeroed() }; + // SAFETY : chemin terminé par NUL, structure de sortie initialisée. + if unsafe { libc::statfs(c.as_ptr(), &mut st) } != 0 { + return None; + } + Some(st) +} + +/// `diskFreeAt` : octets libres pour l'utilisateur, -1 si inconnu. +fn disk_free_at(dir: &str) -> i64 { + #[cfg(unix)] + { + match statfs(dir) { + Some(st) => { + // Types selon la plateforme (u32 ou u64) : les conversions ne + // sont pas toutes superflues. + #[allow(clippy::unnecessary_cast)] + let free = (st.f_bavail as u64).saturating_mul(st.f_bsize as u64); + if free > 1 << 62 { -1 } else { free as i64 } + } + None => -1, + } + } + #[cfg(not(unix))] + { + let _ = dir; + -1 + } +} + +/// `diskFree` : remonte jusqu'au premier parent qui existe. +pub fn disk_free(dir: &str) -> i64 { + let Some(mut d) = go_abs(dir.trim()) else { + return -1; + }; + loop { + if Path::new(&d).is_dir() { + return disk_free_at(&d); + } + let parent = crate::paths::go_dir(&d); + if parent == d { + return -1; + } + d = parent; + } +} + +/// `diskFreeReliable` : l'espace annoncé peut-il servir à REFUSER ? Non sur +/// FUSE (shfs d'Unraid…), overlayfs, NFS, SMB, Ceph, GlusterFS. +pub fn disk_free_reliable(dir: &str) -> bool { + #[cfg(target_os = "linux")] + { + const UNRELIABLE: [u32; 7] = [ + 0x6573_5546, // FUSE + 0x794c_7630, // overlayfs + 0x6969, // NFS + 0x517b, // SMB1 + 0xfe53_4d42, // SMB2+ + 0x00c3_6400, // Ceph + 0x5846_5342, // GlusterFS + ]; + match statfs(dir) { + Some(st) => !UNRELIABLE.contains(&(st.f_type as u32)), + None => false, + } + } + #[cfg(not(target_os = "linux"))] + { + let _ = dir; + true + } +} + +/// `checkDiskSpace` : refus si ça ne tient pas (mesure fiable seulement). +pub fn check_disk_space(dir: &str, size: i64) -> Result<(), String> { + let free = disk_free(dir); + if size <= 0 || free < 0 || !disk_free_reliable(dir) { + return Ok(()); + } + if free < size + DL_SPACE_MARGIN { + return Err(format!( + "espace insuffisant sur {dir} : {} libres, {} nécessaires", + human_bytes(free), + human_bytes(size + DL_SPACE_MARGIN) + )); + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// Les téléchargements +// --------------------------------------------------------------------------- + +/// Un téléchargement (`dlState`), tel que `/status` le rend. +#[derive(Debug, Clone, Default, Serialize, PartialEq)] +pub struct DlState { + pub filename: String, + pub url: String, + pub dir: String, + pub total: i64, + pub done: i64, + pub speed: i64, + pub conns: i64, + pub parts: i64, + pub part: i64, + pub finished: bool, + pub canceled: bool, + #[serde(rename = "error")] + pub err: String, + pub started_at: i64, +} + +struct Entry { + state: Mutex, + cancel: watch::Sender, +} + +impl Entry { + fn with(&self, f: impl FnOnce(&mut DlState) -> R) -> R { + f(&mut self.state.lock().unwrap_or_else(|e| e.into_inner())) + } +} + +/// Une annulation, attendue ou constatée. +#[derive(Clone)] +struct Stop(watch::Receiver); + +impl Stop { + fn is_set(&self) -> bool { + *self.0.borrow() + } + async fn wait(&self) { + let mut rx = self.0.clone(); + let _ = rx.wait_for(|v| *v).await; + } +} + +/// Ce qu'un téléchargement doit savoir de la machine. +#[derive(Clone)] +pub struct DlEnv { + /// Jeton Hugging Face ("" : aucun). + pub hf_token: String, + /// Version de Loki (User-Agent). + pub version: String, +} + +/// Les téléchargements de ce lancement, par nom de fichier. +pub struct Downloads { + map: Mutex>>, + client: reqwest::Client, +} + +impl Default for Downloads { + fn default() -> Self { + Self::new() + } +} + +/// Fin anticipée d'une part : annulation, ou erreur (message). +enum Halt { + Canceled, + Err(String), +} + +impl Downloads { + pub fn new() -> Self { + // HTTP/1.1 : de vraies connexions parallèles, sans fenêtre de flux h2 + // partagée ; pas de délai global (gros fichiers). + crate::chat::engine::install_crypto(); + let client = reqwest::Client::builder() + .http1_only() + .pool_max_idle_per_host(64) + .connect_timeout(Duration::from_secs(20)) + .build() + .unwrap_or_default(); + Downloads { + map: Mutex::new(HashMap::new()), + client, + } + } + + fn map(&self) -> std::sync::MutexGuard<'_, HashMap>> { + self.map.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// `handleModelDownloadStatus` : tous les téléchargements connus. + pub fn status(&self) -> Vec { + let mut v: Vec = self.map().values().map(|e| e.with(|s| s.clone())).collect(); + v.sort_by_key(|s| (s.started_at, s.filename.clone())); + v + } + + /// Annule ; `false` : inconnu. + pub fn cancel(&self, name: &str) -> bool { + match self.map().get(name) { + Some(e) => { + let _ = e.cancel.send(true); + true + } + None => false, + } + } + + fn request(&self, env: &DlEnv, url: &str, range: Option<&str>) -> reqwest::RequestBuilder { + let mut rb = self + .client + .get(url) + .header("User-Agent", format!("loki/{}", env.version)) + // Jamais gzip sur un .gguf : ça casse les plages d'octets. + .header("Accept-Encoding", "identity"); + // Le jeton n'accompagne QUE les liens Hugging Face (reqwest le retire + // quand une redirection change d'hôte). + if !env.hf_token.is_empty() + && reqwest::Url::parse(url) + .ok() + .and_then(|u| u.host_str().map(is_hf_host)) + .unwrap_or(false) + { + rb = rb.header("Authorization", format!("Bearer {}", env.hf_token)); + } + if let Some(r) = range { + rb = rb.header("Range", r); + } + rb + } + + fn source_error(env: &DlEnv, resp: &reqwest::Response, url: &str) -> String { + let code = resp + .headers() + .get("x-error-code") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + dl_source_error(resp.status().as_u16(), code, url, !env.hf_token.is_empty()) + } + + /// `dlProbe` : un octet demandé, pour la taille et le support des plages. + pub async fn probe(&self, env: &DlEnv, url: &str) -> Result<(i64, bool), String> { + let resp = self + .request(env, url, Some("bytes=0-0")) + .send() + .await + .map_err(|e| e.to_string())?; + match resp.status().as_u16() { + 206 => Ok(resp + .headers() + .get("content-range") + .and_then(|v| v.to_str().ok()) + .map(content_range_total) + .filter(|t| *t > 0) + .map_or((0, false), |t| (t, true))), + 200 => Ok((resp.content_length().map_or(-1, |n| n as i64), false)), + _ => Err(Self::source_error(env, &resp, url)), + } + } + + /// `handleModelDownload` : lance en arrière-plan ; rend la réponse. + pub fn start( + self: &Arc, + env: DlEnv, + dirs: &ModelDirs, + url: &str, + dir: &str, + ) -> (StatusCode, Value) { + let bad = |e: String| (StatusCode::BAD_REQUEST, json!({ "ok": false, "error": e })); + let (dl_url, name) = match normalize_hf_url(url) { + Ok(v) => v, + Err(e) => return bad(e), + }; + let (urls, names) = shard_url_set(&dl_url, &name); + let name = names[0].clone(); + let mut dests = Vec::new(); + for n in &names { + match dirs.download_dest_path(n, dir) { + Ok(d) => dests.push(d), + Err(e) => return bad(e), + } + } + let dest_dir = crate::paths::go_dir(&dests[0]); + if let Err(e) = std::fs::create_dir_all(&dest_dir) { + return bad(format!( + "dossier de destination inaccessible : {}", + crate::paths::os_err("mkdir", &dest_dir, &e) + )); + } + // Tranches déjà là : pas retéléchargées. + let todo: Vec<(String, String)> = urls + .into_iter() + .zip(dests.iter().cloned()) + .filter(|(_, d)| !present(d)) + .collect(); + let mut map = self.map(); + if map.get(&name).is_some_and(|e| !e.with(|s| s.finished)) { + return ( + StatusCode::CONFLICT, + json!({ "ok": false, "error": format!("téléchargement déjà en cours pour {name}") }), + ); + } + if todo.is_empty() { + drop(map); + // Rien à faire : le modèle est installé ; l'UI le sélectionne. + let value = dirs + .model_presence(&name, dir) + .map(|(v, _)| v) + .unwrap_or_else(|| dirs.model_picker_value(&dest_dir, &name)); + return ( + StatusCode::OK, + json!({ "ok": true, "filename": name, "exists": true, "value": value, "parts": 0 }), + ); + } + let (tx, rx) = watch::channel(false); + let entry = Arc::new(Entry { + state: Mutex::new(DlState { + filename: name.clone(), + url: todo[0].0.clone(), + dir: dest_dir, + parts: todo.len() as i64, + started_at: crate::chat::persist::now_unix(), + ..Default::default() + }), + cancel: tx, + }); + map.insert(name.clone(), entry.clone()); + drop(map); + let parts = todo.len(); + let me = self.clone(); + tokio::spawn(async move { me.run_set(env, entry, Stop(rx), todo).await }); + ( + StatusCode::OK, + json!({ "ok": true, "filename": name, "parts": parts }), + ) + } + + /// `runDownloadSet` : les tranches l'une après l'autre ; progression sur + /// l'ensemble ; annulation ou échec : tout ce que CE téléchargement a + /// écrit est effacé. + async fn run_set( + self: Arc, + env: DlEnv, + entry: Arc, + stop: Stop, + todo: Vec<(String, String)>, + ) { + let done = Arc::new(AtomicI64::new(0)); + let finish = |res: Result<(), Halt>| { + entry.with(|s| { + match &res { + Err(Halt::Canceled) => { + s.canceled = true; + s.speed = 0; + } + _ if stop.is_set() => { + s.canceled = true; + s.speed = 0; + } + Err(Halt::Err(e)) => { + s.err = e.clone(); + s.speed = 0; + } + Ok(()) => s.done = done.load(Ordering::Relaxed), + } + s.finished = true; + }); + }; + // Sonde de TOUTES les tranches avant d'écrire quoi que ce soit. + let mut sizes = Vec::new(); + let mut grand = 0i64; + for (u, _) in &todo { + let r = tokio::select! { + r = self.probe(&env, u) => r.map_err(Halt::Err), + _ = stop.wait() => Err(Halt::Canceled), + }; + match r { + Ok((n, ranged)) => { + grand += n; + sizes.push((n, ranged)); + } + Err(h) => return finish(Err(h)), + } + } + let dir = entry.with(|s| s.dir.clone()); + if let Err(e) = check_disk_space(&dir, grand) { + return finish(Err(Halt::Err(e))); + } + entry.with(|s| { + s.total = grand; + s.parts = todo.len() as i64; + }); + // Progression et vitesse lissée, une fois par seconde. + let ticker = { + let (entry, done) = (entry.clone(), done.clone()); + tokio::spawn(async move { + let mut t = tokio::time::interval(Duration::from_secs(1)); + t.tick().await; + let (mut last, mut last_at) = (0i64, tokio::time::Instant::now()); + loop { + let now = t.tick().await; + let cur = done.load(Ordering::Relaxed); + let dt = now.duration_since(last_at).as_secs_f64(); + entry.with(|s| { + s.done = cur; + if dt > 0.0 { + let inst = ((cur - last) as f64 / dt) as i64; + s.speed = if s.speed == 0 { + inst + } else { + (s.speed * 2 + inst) / 3 + }; + } + }); + (last, last_at) = (cur, now); + } + }) + }; + let mut written: Vec = Vec::new(); + for (i, ((u, dest), (total, ranged))) in todo.iter().zip(sizes).enumerate() { + entry.with(|s| s.part = i as i64 + 1); + if let Err(h) = self + .one_part(&env, &entry, &stop, u, dest, total, ranged, &done) + .await + { + ticker.abort(); + for p in &written { + let _ = std::fs::remove_file(p); + } + return finish(Err(h)); + } + written.push(dest.clone()); + } + ticker.abort(); + finish(Ok(())); + } + + /// `dlOnePart` : un fichier dans un `.part`, renommé au succès. + #[allow(clippy::too_many_arguments)] + async fn one_part( + &self, + env: &DlEnv, + entry: &Entry, + stop: &Stop, + url: &str, + dest: &str, + total: i64, + ranged: bool, + done: &Arc, + ) -> Result<(), Halt> { + let tmp = format!("{dest}.part"); + let mut conns = 1usize; + if ranged && total > 0 { + let max = ((total + DL_MIN_CHUNK - 1) / DL_MIN_CHUNK) as usize; + conns = dl_conns().min(max).max(1); + } + entry.with(|s| s.conns = conns as i64); + let file = std::fs::File::create(&tmp) + .map_err(|e| Halt::Err(crate::paths::os_err("open", &tmp, &e)))?; + if conns > 1 + && let Err(e) = file.set_len(total as u64) + { + let _ = std::fs::remove_file(&tmp); + return Err(Halt::Err(crate::paths::os_err("truncate", &tmp, &e))); + } + let file = Arc::new(PartFile { + f: file, + path: tmp.clone(), + }); + let res = if conns <= 1 { + self.chunk(env, stop, &file, url, 0, total - 1, total <= 0, done) + .await + } else { + let size = total / conns as i64; + let parts = (0..conns).map(|i| { + let start = i as i64 * size; + let end = if i == conns - 1 { + total - 1 + } else { + start + size - 1 + }; + self.chunk(env, stop, &file, url, start, end, false, done) + }); + futures_util::future::join_all(parts) + .await + .into_iter() + .find(Result::is_err) + .unwrap_or(Ok(())) + }; + if let Err(h) = res { + drop(file); + let _ = std::fs::remove_file(&tmp); + return Err(h); + } + drop(file); + std::fs::rename(&tmp, dest).map_err(|e| { + let _ = std::fs::remove_file(&tmp); + Halt::Err(format!( + "rename {tmp} {dest}: {}", + crate::tools::files::os_err_text(&e) + )) + }) + } + + /// `dlChunk` : [start, end] écrit en place, reprise où la connexion a + /// lâché ; `whole` : tout le corps d'un flux (source sans plages). + #[allow(clippy::too_many_arguments)] + async fn chunk( + &self, + env: &DlEnv, + stop: &Stop, + file: &Arc, + url: &str, + start: i64, + end: i64, + whole: bool, + done: &Arc, + ) -> Result<(), Halt> { + let mut pos = start; + let mut last_err = String::new(); + for attempt in 0..DL_ATTEMPTS { + if stop.is_set() { + return Err(Halt::Canceled); + } + if attempt > 0 { + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(attempt)) => {} + _ = stop.wait() => return Err(Halt::Canceled), + } + } + let range = if !whole { + if pos > end { + return Ok(()); + } + Some(format!("bytes={pos}-{end}")) + } else if pos > start { + Some(format!("bytes={pos}-")) // reprise au mieux + } else { + None + }; + let resp = tokio::select! { + r = self.request(env, url, range.as_deref()).send() => r, + _ = stop.wait() => return Err(Halt::Canceled), + }; + let resp = match resp { + Ok(r) => r, + Err(e) => { + last_err = e.to_string(); + continue; + } + }; + let status = resp.status().as_u16(); + if status != 200 && status != 206 { + return Err(Halt::Err(Self::source_error(env, &resp, url))); + } + if status == 200 && pos > start { + // Reprise refusée : le corps repart de zéro. + done.fetch_add(start - pos, Ordering::Relaxed); + pos = start; + } + match copy_at(stop, file, resp, pos, done).await { + (n, Ok(())) => { + let _ = n; + return Ok(()); + } + (_, Err(Halt::Canceled)) => return Err(Halt::Canceled), + (n, Err(Halt::Err(e))) => { + pos += n; + last_err = e; + } + } + } + Err(Halt::Err(last_err)) + } +} + +/// `dlCopy` : le corps écrit à partir de `off` ; rend les octets écrits, +/// même en cas d'erreur (pour reprendre). +async fn copy_at( + stop: &Stop, + file: &Arc, + resp: reqwest::Response, + off: i64, + done: &Arc, +) -> (i64, Result<(), Halt>) { + let mut stream = resp.bytes_stream(); + let mut written = 0i64; + let mut buf: Vec = Vec::with_capacity(DL_WRITE_BUF); + loop { + let next = tokio::select! { + n = stream.next() => n, + _ = stop.wait() => return (written, Err(Halt::Canceled)), + }; + let eof = next.is_none(); + match next { + Some(Ok(b)) => buf.extend_from_slice(&b), + Some(Err(e)) => { + let n = flush(file, &mut buf, off + written).await; + return match n { + Ok(k) => { + written += k; + done.fetch_add(k, Ordering::Relaxed); + (written, Err(Halt::Err(e.to_string()))) + } + Err(we) => (written, Err(Halt::Err(we))), + }; + } + None => {} + } + if buf.len() >= DL_WRITE_BUF || (eof && !buf.is_empty()) { + match flush(file, &mut buf, off + written).await { + Ok(k) => { + written += k; + done.fetch_add(k, Ordering::Relaxed); + } + Err(e) => return (written, Err(Halt::Err(e))), + } + } + if eof { + return (written, Ok(())); + } + } +} + +/// Écrit le tampon à sa place, hors du runtime ; le vide. +async fn flush(file: &Arc, buf: &mut Vec, at: i64) -> Result { + let data = std::mem::take(buf); + let n = data.len() as i64; + let f = file.clone(); + tokio::task::spawn_blocking(move || { + write_all_at(&f.f, &data, at as u64).map_err(|e| crate::paths::os_err("write", &f.path, &e)) + }) + .await + .map_err(|e| e.to_string())??; + Ok(n) +} + +/// Le `.part` en cours d'écriture, et son chemin (messages d'erreur). +struct PartFile { + f: std::fs::File, + path: String, +} + +#[cfg(unix)] +fn write_all_at(f: &std::fs::File, data: &[u8], at: u64) -> std::io::Result<()> { + use std::os::unix::fs::FileExt; + f.write_all_at(data, at) +} + +#[cfg(windows)] +fn write_all_at(f: &std::fs::File, mut data: &[u8], mut at: u64) -> std::io::Result<()> { + use std::os::windows::fs::FileExt; + while !data.is_empty() { + let n = f.seek_write(data, at)?; + data = &data[n..]; + at += n as u64; + } + Ok(()) +} + +/// `cleanStalePartFiles` : au démarrage, les `.part` d'un téléchargement tué +/// (crash, redémarrage) ne reprendront jamais : effacés. +pub fn clean_stale_part_files(dirs: &[String]) { + for d in dirs { + let Ok(rd) = std::fs::read_dir(d) else { + continue; + }; + for e in rd.filter_map(Result::ok) { + let name = e.file_name().to_string_lossy().into_owned(); + // `filepath.Glob("*.gguf.part")` : sensible à la casse. + if name.ends_with(".gguf.part") && std::fs::remove_file(e.path()).is_ok() { + println!( + "[models] téléchargement incomplet supprimé : {}", + go_clean(&e.path().to_string_lossy()) + ); + } + } + } +} + +// --------------------------------------------------------------------------- +// Routes +// --------------------------------------------------------------------------- + +/// Les dossiers de modèles et l'environnement du téléchargement, lus en base. +async fn dirs_and_env(ctx: &Ctx, version: String) -> (ModelDirs, DlEnv) { + let stored = gguf::parse_stored_model_dirs( + ctx.store + .get(BK_STATE, gguf::STATE_MODEL_DIRS_KEY) + .await + .ok() + .flatten() + .as_deref(), + ); + let dirs = ModelDirs { + home: ctx.engine.home.clone(), + stored, + env: std::env::var(gguf::MODEL_DIRS_ENV).ok(), + }; + ( + dirs, + DlEnv { + hf_token: hf_token(&ctx.store).await, + version, + }, + ) +} + +/// `hfToken` : le jeton enregistré d'abord, `HF_TOKEN` ensuite. +pub async fn hf_token(store: &Store) -> String { + let t = store.get_str(BK_STATE, "hf_token").await; + if !t.trim().is_empty() { + return t.trim().to_string(); + } + std::env::var("HF_TOKEN") + .unwrap_or_default() + .trim() + .to_string() +} + +async fn body_json(body: Incoming) -> Result { + let b = Limited::new(body, 1 << 20) + .collect() + .await + .map_err(|e| e.to_string())? + .to_bytes(); + if b.iter().all(u8::is_ascii_whitespace) { + return Err("EOF".into()); + } + serde_json::from_slice(&b).map_err(|e| e.to_string()) +} + +fn str_of<'a>(v: &'a Value, k: &str) -> &'a str { + v.get(k).and_then(Value::as_str).unwrap_or("") +} + +/// `from_go` : remise par le Go (tunnel), déjà authentifiée par lui. +pub async fn serve( + ctx: &Ctx, + dl: &Arc, + req: Request, + from_go: bool, +) -> Response { + let (parts, body) = req.into_parts(); + // La version du Go (User-Agent) vient avec l'autorisation ; une remise du + // Go la redemande. + let version = if from_go { + ctx.internal + .auth() + .await + .map(|i| i.version) + .unwrap_or_default() + } else { + match ctx.guard(&parts.headers).await { + Ok(v) => v, + Err(refus) => return *refus, + } + }; + let bad = |e: String| send_json(StatusCode::BAD_REQUEST, &json!({ "ok": false, "error": e })); + match parts.uri.path() { + "/api/models/download/status" => send_json( + StatusCode::OK, + &serde_json::to_value(dl.status()).unwrap_or_default(), + ), + "/api/models/download/cancel" => { + let v = match body_json(body).await { + Ok(v) => v, + Err(e) => return bad(e), + }; + let name = filepath_base(str_of(&v, "filename").trim()); + if dl.cancel(&name) { + send_json(StatusCode::OK, &json!({ "ok": true })) + } else { + send_json( + StatusCode::NOT_FOUND, + &json!({ "ok": false, "error": format!("aucun téléchargement pour {name}") }), + ) + } + } + "/api/models/download" => { + let v = match body_json(body).await { + Ok(v) => v, + Err(e) => return bad(e), + }; + let (dirs, env) = dirs_and_env(ctx, version).await; + let (url, dir) = (str_of(&v, "url").to_string(), str_of(&v, "dir").to_string()); + let dl = dl.clone(); + // Disque (dossiers, présence, mkdir) hors du runtime ; le + // téléchargement lui-même y repart (tokio::spawn). + let (code, out) = tokio::task::spawn_blocking(move || dl.start(env, &dirs, &url, &dir)) + .await + .unwrap_or_else(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + json!({ "ok": false, "error": e.to_string() }), + ) + }); + send_json(code, &out) + } + "/api/models/download/probe" => { + let v = match body_json(body).await { + Ok(v) => v, + Err(e) => return bad(e), + }; + let (dirs, env) = dirs_and_env(ctx, version).await; + let (dl_url, name) = match normalize_hf_url(str_of(&v, "url")) { + Ok(x) => x, + Err(e) => return bad(e), + }; + let dir = match dirs.resolve_download_dir(str_of(&v, "dir")) { + Ok(d) => d, + Err(e) => return bad(e), + }; + let (urls, names) = shard_url_set(&dl_url, &name); + let name = names[0].clone(); + // Déjà là (dans N'IMPORTE quel dossier déclaré) : rien à sonder, + // un modèle à sélectionner. + if let Some((value, path)) = dirs.model_presence(&name, "") { + return send_json( + StatusCode::OK, + &json!({ + "ok": true, "filename": name, "dir": crate::paths::go_dir(&path), + "exists": true, "value": value, "size": model_family_size(&path), + "free": disk_free(&dir), "free_exact": disk_free_reliable(&dir), + "enough": true, "parts": urls.len(), + }), + ); + } + let mut total = 0i64; + let probe_all = async { + for u in &urls { + let (n, _) = dl.probe(&env, u).await?; + total += n; + } + Ok::<(), String>(()) + }; + match tokio::time::timeout(Duration::from_secs(30), probe_all).await { + Ok(Ok(())) => {} + Ok(Err(e)) => return bad(e), + Err(_) => return bad("context deadline exceeded".into()), + } + let mut out = json!({ + "ok": true, "filename": name, "dir": dir, "size": total, + "free": disk_free(&dir), "free_exact": disk_free_reliable(&dir), + "enough": true, "parts": urls.len(), + }); + if let Err(e) = check_disk_space(&dir, total) { + out["enough"] = json!(false); + out["error"] = json!(e); + } + send_json(StatusCode::OK, &out) + } + _ => send_json( + StatusCode::NOT_FOUND, + &json!({ "error": "404 page not found" }), + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn noms_et_tailles() { + assert_eq!(url_path_base("/a/b/"), "b"); + assert_eq!(url_path_base(""), "."); + assert_eq!(url_path_base("///"), "/"); + assert_eq!(percent_decode("a%20b%zz%4"), "a b%zz%4"); + assert_eq!(human_bytes(3 << 30), "3.0 Go"); + assert_eq!(human_bytes(5 << 20), "5 Mo"); + assert_eq!(human_bytes(12), "12 o"); + assert_eq!(filepath_base(""), "."); + } + + #[test] + fn espace_disque() { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join("pas/encore/la"); + let free = disk_free(&p.to_string_lossy()); + assert!(free != 0, "{free}"); + assert!(check_disk_space(&d.path().to_string_lossy(), 1).is_ok()); + if free > 0 && disk_free_reliable(&d.path().to_string_lossy()) { + let e = check_disk_space(&d.path().to_string_lossy(), free).unwrap_err(); + assert!(e.starts_with("espace insuffisant sur "), "{e}"); + } + } + + #[test] + fn part_perimes_effaces() { + let d = tempfile::tempdir().unwrap(); + std::fs::write(d.path().join("m.gguf.part"), b"x").unwrap(); + std::fs::write(d.path().join("m.gguf"), b"x").unwrap(); + clean_stale_part_files(&[d.path().to_string_lossy().into_owned()]); + assert!(!d.path().join("m.gguf.part").exists() && d.path().join("m.gguf").exists()); + } +} diff --git a/rust/loki-rs/src/lib.rs b/rust/loki-rs/src/lib.rs index be2668b..e7ab18a 100644 --- a/rust/loki-rs/src/lib.rs +++ b/rust/loki-rs/src/lib.rs @@ -15,6 +15,7 @@ pub mod auth; pub mod chat; pub mod chatapi; pub mod child; +pub mod downloads; pub mod engine; pub mod engstatus; pub mod gguf; diff --git a/rust/loki-rs/src/main.rs b/rust/loki-rs/src/main.rs index 1d735ca..24b75d6 100644 --- a/rust/loki-rs/src/main.rs +++ b/rust/loki-rs/src/main.rs @@ -118,9 +118,17 @@ fn serve() -> ExitCode { // sait (il remet alors à la façade le trafic de chat du tunnel) et // connaît son adresse locale. let chat_rust = chat_held_by_facade(); + // LOKI_DOWNLOADS=rust : les téléchargements de modèles aussi (le Go + // remet alors leurs routes du tunnel et ne vide plus les .part). + let downloads_rust = downloads_held_by_facade(); let mut env = vec![(auth::INTERNAL_ENV.to_string(), token.clone())]; if chat_rust { env.push(("LOKI_CHAT".into(), "rust".into())); + } + if downloads_rust { + env.push(("LOKI_DOWNLOADS".into(), "rust".into())); + } + if chat_rust || downloads_rust { env.push(("LOKI_FACADE_ADDR".into(), format!("127.0.0.1:{web_port}"))); } let sup = child::Supervisor::spawn(go_bin, go_port, env); @@ -153,6 +161,11 @@ fn serve() -> ExitCode { } else { None }; + let downloads = if downloads_rust { + Some(start_downloads(backend, &token, &engine_paths.home).await) + } else { + None + }; let addr = SocketAddr::from(([0, 0, 0, 0], web_port)); let listener = match tokio::net::TcpListener::bind(addr).await { @@ -170,6 +183,7 @@ fn serve() -> ExitCode { host_env: auth::HostEnv::from_env(), engine: engine_paths, chat, + downloads, }; let res = proxy::serve(listener, cfg, shutdown_signal()).await; eprintln!("[loki-rs] arrêt demandé — arrêt du serveur Go"); @@ -189,6 +203,37 @@ fn chat_held_by_facade() -> bool { std::env::var("LOKI_CHAT").is_ok_and(|v| v.trim().eq_ignore_ascii_case("rust")) } +/// `LOKI_DOWNLOADS=rust` : téléchargements de modèles tenus par la façade. +fn downloads_held_by_facade() -> bool { + std::env::var("LOKI_DOWNLOADS").is_ok_and(|v| v.trim().eq_ignore_ascii_case("rust")) +} + +/// Les `.part` d'un téléchargement tué ne reprendront jamais : vidés au +/// démarrage (le Go ne le fait plus quand la façade les tient). +async fn start_downloads( + backend: SocketAddr, + token: &str, + home: &std::path::Path, +) -> std::sync::Arc { + use loki_rs::{downloads, gguf, store}; + let st = store::Store::new(backend, token.to_string()); + let stored = gguf::parse_stored_model_dirs( + st.get(store::BK_STATE, gguf::STATE_MODEL_DIRS_KEY) + .await + .ok() + .flatten() + .as_deref(), + ); + let dirs = gguf::model_dirs( + home, + &stored, + std::env::var(gguf::MODEL_DIRS_ENV).ok().as_deref(), + ); + let _ = tokio::task::spawn_blocking(move || downloads::clean_stale_part_files(&dirs)).await; + eprintln!("[loki-rs] téléchargements de modèles tenus par la façade (LOKI_DOWNLOADS=rust)"); + std::sync::Arc::new(downloads::Downloads::new()) +} + /// Le Go garde-t-il la conversation ? (quelques essais : il vient de démarrer) async fn go_keeps_chat(backend: SocketAddr, token: &str) -> Result, String> { let store = loki_rs::store::Store::new(backend, token.to_string()); diff --git a/rust/loki-rs/src/proxy.rs b/rust/loki-rs/src/proxy.rs index b2b55da..930b561 100644 --- a/rust/loki-rs/src/proxy.rs +++ b/rust/loki-rs/src/proxy.rs @@ -63,6 +63,9 @@ pub struct Config { /// Conversation tenue par la façade (`LOKI_CHAT=rust`) : ses routes sont /// servies ici ; `None` : relayées au Go. pub chat: Option>, + /// Téléchargements de modèles tenus par la façade (`LOKI_DOWNLOADS=rust`) : + /// leurs routes sont servies ici ; `None` : relayées au Go. + pub downloads: Option>, } struct Proxy { @@ -90,6 +93,7 @@ pub async fn serve( host_env: cfg.host_env, engine: cfg.engine, chat: cfg.chat, + downloads: cfg.downloads, internal_token: cfg.internal_token.clone(), }, }); @@ -129,8 +133,10 @@ pub async fn serve( impl Proxy { async fn handle(&self, mut req: Request) -> Response { // Le Go (et lui seul : jeton tiré à chaque démarrage) remet à la - // façade le trafic de chat qui lui arrive sans passer par elle. - let from_go = self.ctx.chat.is_some() && self.ctx.from_go(req.headers()); + // façade le trafic de chat et de téléchargements qui lui arrive sans + // passer par elle. + let from_go = (self.ctx.chat.is_some() || self.ctx.downloads.is_some()) + && self.ctx.from_go(req.headers()); // Jamais de confiance accordée sur la foi d'un en-tête du client : // ni « je viens du tunnel », ni « je suis la façade ». req.headers_mut().remove(VIA_TUNNEL); @@ -163,6 +169,11 @@ impl Proxy { { return chatapi::serve(&self.ctx, chat, req, from_go).await; } + if let Some(dl) = &self.ctx.downloads + && crate::downloads::handles(path) + { + return crate::downloads::serve(&self.ctx, dl, req, from_go).await; + } self.forward(req).await } diff --git a/rust/loki-rs/tests/common/fakego.rs b/rust/loki-rs/tests/common/fakego.rs index 19b92dd..4ef5359 100644 --- a/rust/loki-rs/tests/common/fakego.rs +++ b/rust/loki-rs/tests/common/fakego.rs @@ -427,7 +427,16 @@ async fn handle(st: Arc, req: Request) -> Response /// Lance le faux Go puis la façade devant lui ; rend (adresse de la façade, /// état du faux Go, garde d'arrêt). pub async fn start(home: std::path::PathBuf) -> (SocketAddr, Arc, oneshot::Sender<()>) { - let (f, st, tx, _) = start_opts(home, false, None).await; + let (f, st, tx, _) = start_opts(home, false, None, false).await; + (f, st, tx) +} + +/// Comme [`start`], avec les téléchargements de modèles tenus par la façade +/// (`LOKI_DOWNLOADS=rust`). +pub async fn start_downloads( + home: std::path::PathBuf, +) -> (SocketAddr, Arc, oneshot::Sender<()>) { + let (f, st, tx, _) = start_opts(home, false, None, true).await; (f, st, tx) } @@ -442,7 +451,7 @@ pub async fn start_chat( oneshot::Sender<()>, Arc, ) { - let (f, st, tx, chat) = start_opts(home, true, Some(engine_port)).await; + let (f, st, tx, chat) = start_opts(home, true, Some(engine_port), false).await; (f, st, tx, chat.unwrap()) } @@ -450,6 +459,7 @@ async fn start_opts( home: std::path::PathBuf, with_chat: bool, engine_port: Option, + with_downloads: bool, ) -> ( SocketAddr, Arc, @@ -506,6 +516,7 @@ async fn start_opts( service: "loki-engine".into(), }, chat: chat.clone(), + downloads: with_downloads.then(|| Arc::new(loki_rs::downloads::Downloads::new())), }; tokio::spawn(loki_rs::proxy::serve(fl, cfg, async move { let _ = rx.await; @@ -513,6 +524,32 @@ async fn start_opts( (facade, st, tx, chat) } +/// Requête via la façade (corps JSON) : (code, corps). +pub async fn send( + facade: SocketAddr, + method: &str, + path: &str, + headers: &[(&str, &str)], + body: &str, +) -> (StatusCode, Bytes) { + use hyper_util::client::legacy::Client; + use hyper_util::rt::TokioExecutor; + let mut req = Request::builder() + .method(method) + .uri(format!("http://{facade}{path}")) + .header("content-type", "application/json"); + for (k, v) in headers { + req = req.header(*k, *v); + } + let client: Client<_, Full> = Client::builder(TokioExecutor::new()).build_http(); + let r = client + .request(req.body(Full::new(Bytes::from(body.to_string()))).unwrap()) + .await + .unwrap(); + let (parts, body) = r.into_parts(); + (parts.status, body.collect().await.unwrap().to_bytes()) +} + /// GET via la façade : (code, en-têtes, corps). pub async fn get( facade: SocketAddr, diff --git a/rust/loki-rs/tests/downloads.rs b/rust/loki-rs/tests/downloads.rs new file mode 100644 index 0000000..17e4aad --- /dev/null +++ b/rust/loki-rs/tests/downloads.rs @@ -0,0 +1,692 @@ +//! Lot « modèles » : téléchargements tenus par la façade (`LOKI_DOWNLOADS=rust`). +//! +//! - les fonctions pures rejouent le journal tiré du Go +//! (`internal/loki/parity_dl_test.go` → `tests/parity_models/download.json`) ; +//! - le téléchargement de bout en bout tourne contre un faux Hugging Face +//! (plages d'octets, tranches, dépôt verrouillé, annulation). + +mod common; + +use std::collections::HashMap; +use std::convert::Infallible; +use std::net::SocketAddr; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use bytes::Bytes; +use common::fakego; +use http_body_util::{BodyExt, Full, StreamBody, combinators::BoxBody}; +use hyper::body::{Frame, Incoming}; +use hyper::{Request, Response, StatusCode, service::service_fn}; +use hyper_util::rt::TokioIo; +use loki_rs::downloads::{self, DlEnv, DlState, Downloads, ModelDirs}; +use serde_json::Value; + +// --------------------------------------------------------------------------- +// Journal de référence du Go +// --------------------------------------------------------------------------- + +#[test] +fn parite_des_fonctions_pures() { + let fx: Value = + serde_json::from_str(include_str!("parity_models/download.json")).expect("journal"); + let s = |v: &Value, k: &str| v.get(k).and_then(Value::as_str).unwrap_or("").to_string(); + let list = |v: &Value, k: &str| -> Vec { + v.get(k) + .and_then(Value::as_array) + .map(|a| a.iter().map(|x| x.as_str().unwrap().to_string()).collect()) + .unwrap_or_default() + }; + + for c in fx["normalize"].as_array().unwrap() { + let input = s(c, "in"); + match downloads::normalize_hf_url(&input) { + Err(e) => assert_eq!(e, s(c, "err"), "normalize({input:?})"), + Ok((url, name)) => { + assert_eq!(s(c, "err"), "", "normalize({input:?}) aurait dû échouer"); + assert_eq!(url, s(c, "url"), "url de {input:?}"); + assert_eq!(name, s(c, "name"), "nom de {input:?}"); + let (urls, names) = downloads::shard_url_set(&url, &name); + assert_eq!(urls, list(c, "urls"), "tranches de {input:?}"); + assert_eq!(names, list(c, "names"), "noms de {input:?}"); + } + } + } + for (v, want) in fx["content_range"].as_object().unwrap() { + assert_eq!( + downloads::content_range_total(v), + want.as_i64().unwrap(), + "{v:?}" + ); + } + for c in fx["source"].as_array().unwrap() { + let got = downloads::dl_source_error( + c["status"].as_u64().unwrap() as u16, + &s(c, "code"), + &s(c, "url"), + c["token_set"].as_bool().unwrap(), + ); + assert_eq!(got, s(c, "err"), "{c}"); + } + for (h, want) in fx["hf_host"].as_object().unwrap() { + assert_eq!(downloads::is_hf_host(h), want.as_bool().unwrap(), "{h}"); + } + for (u, want) in fx["hf_repo"].as_object().unwrap() { + assert_eq!( + downloads::hf_repo_from_url(u), + want.as_str().unwrap(), + "{u}" + ); + } + + let tmp = tempfile::tempdir().unwrap(); + let home = tmp.path().canonicalize().unwrap(); + let home_s = home.to_string_lossy().into_owned(); + std::fs::create_dir_all(home.join("disque")).unwrap(); + let dirs = ModelDirs { + home: home.clone(), + stored: vec![format!("{home_s}/disque")], + env: None, + }; + for c in fx["dest"].as_array().unwrap() { + let dir = s(c, "dir").replace("{{HOME}}", &home_s); + let got = match dirs.download_dest_path(&s(c, "name"), &dir) { + Ok(p) => ("path", p.replace(&home_s, "{{HOME}}")), + Err(e) => ("err", e.replace(&home_s, "{{HOME}}")), + }; + let want = if c.get("path").is_some() { + ("path", s(c, "path")) + } else { + ("err", s(c, "err")) + }; + assert_eq!(got, want, "{c}"); + } +} + +#[test] +fn tailles_lisibles_comme_le_go() { + assert_eq!(downloads::human_bytes(0), "0 o"); + assert_eq!(downloads::human_bytes((1 << 20) - 1), "1048575 o"); + assert_eq!(downloads::human_bytes(1 << 20), "1 Mo"); + assert_eq!(downloads::human_bytes(1536 << 20), "1.5 Go"); +} + +// --------------------------------------------------------------------------- +// Faux Hugging Face +// --------------------------------------------------------------------------- + +type Body = BoxBody; + +#[derive(Default)] +struct FakeHf { + files: Mutex>>>, + /// Chemin → (statut, X-Error-Code) imposés. + refuse: Mutex>, + /// Ignore les plages (200, corps entier). + no_ranges: Mutex, + /// Corps envoyé au goutte-à-goutte (annulation). + slow: Mutex, + /// « chemin Range » de chaque requête, et en-têtes Authorization vus. + log: Mutex>, + auth: Mutex>, +} + +fn full(status: u16, body: impl Into) -> Response { + let mut r = Response::new(Full::new(body.into()).map_err(|e| match e {}).boxed()); + *r.status_mut() = StatusCode::from_u16(status).unwrap(); + r +} + +async fn hf_handle(hf: Arc, req: Request) -> Response { + let path = req.uri().path().to_string(); + let range = req + .headers() + .get("range") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); + if let Some(a) = req.headers().get("authorization") { + hf.auth + .lock() + .unwrap() + .push(a.to_str().unwrap().to_string()); + } + hf.log.lock().unwrap().push(format!("{path} {range}")); + if let Some((st, code)) = hf.refuse.lock().unwrap().get(&path).cloned() { + let mut r = full(st, "refus"); + if !code.is_empty() { + r.headers_mut() + .insert("x-error-code", code.parse().unwrap()); + } + return r; + } + let Some(data) = hf.files.lock().unwrap().get(&path).cloned() else { + return full(404, "absent"); + }; + let total = data.len(); + let (status, start, end) = match range.strip_prefix("bytes=") { + Some(r) if !*hf.no_ranges.lock().unwrap() => { + let (a, b) = r.split_once('-').unwrap(); + let a: usize = a.parse().unwrap(); + let b: usize = if b.is_empty() { + total - 1 + } else { + b.parse::().unwrap().min(total - 1) + }; + (206, a, b) + } + _ => (200, 0, total - 1), + }; + let chunk = Bytes::copy_from_slice(&data[start..=end]); + let mut r = if *hf.slow.lock().unwrap() && chunk.len() > 1 { + // 64 Kio toutes les 20 ms : le téléchargement dure, on peut l'annuler. + let s = futures_util::stream::unfold(chunk, |mut rest| async move { + if rest.is_empty() { + return None; + } + tokio::time::sleep(Duration::from_millis(20)).await; + let part = rest.split_to(rest.len().min(64 << 10)); + Some((Ok::<_, Infallible>(Frame::data(part)), rest)) + }); + let mut r = Response::new(StreamBody::new(s).boxed()); + *r.status_mut() = StatusCode::from_u16(status).unwrap(); + r + } else { + full(status, chunk) + }; + if status == 206 { + r.headers_mut().insert( + "content-range", + format!("bytes {start}-{end}/{total}").parse().unwrap(), + ); + } + r +} + +async fn fake_hf() -> (SocketAddr, Arc) { + let hf = Arc::new(FakeHf::default()); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + let h = hf.clone(); + tokio::spawn(async move { + while let Ok((s, _)) = l.accept().await { + let h = h.clone(); + tokio::spawn(async move { + let svc = service_fn(move |req| { + let h = h.clone(); + async move { Ok::<_, Infallible>(hf_handle(h, req).await) } + }); + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection(TokioIo::new(s), svc) + .await; + }); + } + }); + (addr, hf) +} + +/// Contenu reconnaissable : un décalage d'octets se verrait. +fn payload(n: usize, seed: u8) -> Vec { + (0..n) + .map(|i| (i as u32).wrapping_mul(2_654_435_761).to_le_bytes()[1] ^ seed) + .collect() +} + +fn env() -> DlEnv { + DlEnv { + hf_token: "hf_secret".into(), + version: "9.9.9".into(), + } +} + +fn dirs_at(home: &Path) -> ModelDirs { + ModelDirs { + home: home.to_path_buf(), + stored: vec![], + env: None, + } +} + +async fn wait_finished(dl: &Downloads, name: &str) -> DlState { + for _ in 0..1000 { + if let Some(s) = dl.status().into_iter().find(|s| s.filename == name) + && s.finished + { + return s; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("téléchargement {name} jamais terminé : {:?}", dl.status()); +} + +fn models(home: &Path) -> PathBuf { + home.join("models") +} + +fn entries(dir: &Path) -> Vec { + let mut v: Vec = std::fs::read_dir(dir) + .map(|rd| { + rd.filter_map(Result::ok) + .map(|e| e.file_name().to_string_lossy().into_owned()) + .collect() + }) + .unwrap_or_default(); + v.sort(); + v +} + +// --------------------------------------------------------------------------- +// De bout en bout +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn un_fichier_sur_plusieurs_connexions() { + let (addr, hf) = fake_hf().await; + let data = payload(40 << 20, 7); // 40 Mio → 3 parts de 16 Mio au plus + hf.files + .lock() + .unwrap() + .insert("/u/r/resolve/main/m.gguf".into(), Arc::new(data.clone())); + let tmp = tempfile::tempdir().unwrap(); + let dl = Arc::new(Downloads::new()); + let (code, out) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/u/r/resolve/main/m.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK, "{out}"); + assert_eq!(out["filename"], "m.gguf"); + assert_eq!(out["parts"], 1); + let st = wait_finished(&dl, "m.gguf").await; + assert_eq!(st.err, ""); + assert!(!st.canceled); + assert_eq!( + (st.total, st.done, st.conns, st.parts, st.part), + (40 << 20, 40 << 20, 3, 1, 1) + ); + assert_eq!(st.dir, models(tmp.path()).to_string_lossy()); + let got = std::fs::read(models(tmp.path()).join("m.gguf")).unwrap(); + assert!(got == data, "contenu différent"); + assert_eq!(entries(&models(tmp.path())), ["m.gguf"]); + // Sonde d'un octet, puis trois plages disjointes qui couvrent le fichier. + let log = hf.log.lock().unwrap().clone(); + assert_eq!(log[0], "/u/r/resolve/main/m.gguf bytes=0-0"); + let mut ranges: Vec<&str> = log[1..] + .iter() + .map(|l| l.split(' ').nth(1).unwrap()) + .collect(); + ranges.sort(); + assert_eq!( + ranges, + [ + "bytes=0-13981012", + "bytes=13981013-27962025", + "bytes=27962026-41943039" + ] + ); + // Hôte hors Hugging Face : le jeton ne part pas. + assert!(hf.auth.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn source_sans_plages_un_seul_flux() { + let (addr, hf) = fake_hf().await; + let data = payload(3 << 20, 1); + hf.files + .lock() + .unwrap() + .insert("/m.gguf".into(), Arc::new(data.clone())); + *hf.no_ranges.lock().unwrap() = true; + let tmp = tempfile::tempdir().unwrap(); + let dl = Arc::new(Downloads::new()); + let (code, _) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/m.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK); + let st = wait_finished(&dl, "m.gguf").await; + assert_eq!( + (st.err.as_str(), st.total, st.done, st.conns), + ("", 3 << 20, 3 << 20, 1) + ); + assert!(std::fs::read(models(tmp.path()).join("m.gguf")).unwrap() == data); +} + +#[tokio::test] +async fn modele_en_tranches_rapatrie_en_entier() { + let (addr, hf) = fake_hf().await; + let parts: Vec> = (0..3).map(|i| payload(1 << 20, i as u8)).collect(); + for (i, p) in parts.iter().enumerate() { + hf.files.lock().unwrap().insert( + format!("/u/r/resolve/main/Q4/m-0000{}-of-00003.gguf", i + 1), + Arc::new(p.clone()), + ); + } + let tmp = tempfile::tempdir().unwrap(); + let m = models(tmp.path()); + std::fs::create_dir_all(&m).unwrap(); + // La première est déjà là : pas retéléchargée. + std::fs::write(m.join("m-00001-of-00003.gguf"), &parts[0]).unwrap(); + let dl = Arc::new(Downloads::new()); + let (code, out) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/u/r/resolve/main/Q4/m-00002-of-00003.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK, "{out}"); + assert_eq!( + (out["filename"].as_str(), out["parts"].as_i64()), + (Some("m-00001-of-00003.gguf"), Some(2)) + ); + let st = wait_finished(&dl, "m-00001-of-00003.gguf").await; + assert_eq!( + (st.err.as_str(), st.total, st.done, st.parts, st.part), + ("", 2 << 20, 2 << 20, 2, 2) + ); + for (i, p) in parts.iter().enumerate() { + let got = std::fs::read(m.join(format!("m-0000{}-of-00003.gguf", i + 1))).unwrap(); + assert!(&got == p, "tranche {}", i + 1); + } + assert!( + !hf.log + .lock() + .unwrap() + .iter() + .any(|l| l.contains("00001-of")), + "tranche présente retéléchargée" + ); + // Tout est là : rien à faire, le modèle est à sélectionner. + let (code, out) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/u/r/resolve/main/Q4/m-00003-of-00003.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK); + assert_eq!(out["exists"], true); + assert_eq!(out["value"], "m-00001-of-00003.gguf"); + assert_eq!(out["parts"], 0); +} + +#[tokio::test] +async fn depot_verrouille_rien_ne_reste() { + let (addr, hf) = fake_hf().await; + for i in 1..=2 { + hf.files.lock().unwrap().insert( + format!("/g/d/resolve/main/m-0000{i}-of-00002.gguf"), + Arc::new(payload(1 << 20, i)), + ); + } + hf.refuse.lock().unwrap().insert( + "/g/d/resolve/main/m-00002-of-00002.gguf".into(), + (401, "GatedRepo".into()), + ); + let tmp = tempfile::tempdir().unwrap(); + let dl = Arc::new(Downloads::new()); + let (code, _) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/g/d/resolve/main/m-00001-of-00002.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK); + let st = wait_finished(&dl, "m-00001-of-00002.gguf").await; + // Refusé dès la sonde : rien n'a été écrit. + assert_eq!( + st.err, + "dépôt à accès restreint : ses conditions doivent être acceptées sur huggingface.co — le jeton Hugging Face enregistré n'y donne pas accès (expiré, ou conditions non acceptées avec ce compte)" + ); + assert_eq!(entries(&models(tmp.path())), Vec::::new()); + assert!( + hf.log + .lock() + .unwrap() + .iter() + .all(|l| l.ends_with("bytes=0-0")) + ); +} + +#[tokio::test] +async fn panne_pendant_la_deuxieme_tranche_efface_la_premiere() { + let (addr, hf) = fake_hf().await; + let tmp = tempfile::tempdir().unwrap(); + let dl = Arc::new(Downloads::new()); + for i in 1..=2 { + hf.files.lock().unwrap().insert( + format!("/m-0000{i}-of-00002.gguf"), + Arc::new(payload(1 << 20, i)), + ); + } + // La sonde passe, le transfert de la deuxième échoue (fichier retiré). + let (code, _) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/m-00001-of-00002.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK); + // Retire la deuxième dès que la première est écrite. + let m = models(tmp.path()); + for _ in 0..500 { + if m.join("m-00001-of-00002.gguf").exists() { + break; + } + if dl.status()[0].finished { + break; + } + tokio::time::sleep(Duration::from_millis(2)).await; + } + let st0 = dl.status()[0].clone(); + if st0.finished { + // Trop rapide pour la machine : le scénario ne se joue pas. + return; + } + hf.files.lock().unwrap().remove("/m-00002-of-00002.gguf"); + let st = wait_finished(&dl, "m-00001-of-00002.gguf").await; + if st.err.is_empty() { + return; // deuxième tranche déjà partie avant le retrait + } + assert_eq!(st.err, "lien introuvable (HTTP 404)"); + assert_eq!(entries(&m), Vec::::new(), "tranches laissées"); +} + +#[tokio::test] +async fn annulation_efface_le_part() { + let (addr, hf) = fake_hf().await; + hf.files + .lock() + .unwrap() + .insert("/m.gguf".into(), Arc::new(payload(8 << 20, 3))); + *hf.slow.lock().unwrap() = true; + let tmp = tempfile::tempdir().unwrap(); + let dl = Arc::new(Downloads::new()); + let (code, _) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/m.gguf"), + "", + ); + assert_eq!(code, StatusCode::OK); + // Deuxième lancement pendant le premier : refusé. + let (code, out) = dl.start( + env(), + &dirs_at(tmp.path()), + &format!("http://{addr}/m.gguf"), + "", + ); + assert_eq!(code, StatusCode::CONFLICT); + assert_eq!(out["error"], "téléchargement déjà en cours pour m.gguf"); + let part = models(tmp.path()).join("m.gguf.part"); + for _ in 0..500 { + if part.exists() { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(part.exists(), "pas de .part en cours"); + assert!(dl.cancel("m.gguf")); + assert!(!dl.cancel("autre.gguf")); + let st = wait_finished(&dl, "m.gguf").await; + assert!(st.canceled && st.err.is_empty(), "{st:?}"); + assert_eq!(st.speed, 0); + assert_eq!(entries(&models(tmp.path())), Vec::::new()); +} + +#[tokio::test] +async fn refus_avant_de_lancer() { + let tmp = tempfile::tempdir().unwrap(); + let dl = Arc::new(Downloads::new()); + let d = dirs_at(tmp.path()); + let (code, out) = dl.start(env(), &d, "pas un lien", ""); + assert_eq!( + (code, out["error"].as_str()), + ( + StatusCode::BAD_REQUEST, + Some("lien invalide (http/https attendu)") + ) + ); + let (code, out) = dl.start(env(), &d, "http://h/m.gguf", "/etc"); + assert_eq!(code, StatusCode::BAD_REQUEST); + assert_eq!( + out["error"], + "dossier non autorisé : /etc — ajoute-le dans « Dossiers de modèles »" + ); + assert!(dl.status().is_empty()); +} + +// --------------------------------------------------------------------------- +// Par la façade +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn routes_servies_par_la_facade() { + let (addr, hf) = fake_hf().await; + hf.files.lock().unwrap().insert( + "/u/r/resolve/main/m.gguf".into(), + Arc::new(payload(1 << 20, 9)), + ); + let tmp = tempfile::tempdir().unwrap(); + let home = tmp.path().canonicalize().unwrap(); + let (facade, go, _stop) = fakego::start_downloads(home.clone()).await; + let url = format!("http://{addr}/u/r/resolve/main/m.gguf"); + let body = format!(r#"{{"url":"{url}","dir":""}}"#); + + // Sonde : taille, dossier, espace. + let (code, b) = fakego::send(facade, "POST", "/api/models/download/probe", &[], &body).await; + assert_eq!(code, StatusCode::OK, "{b:?}"); + let v: Value = serde_json::from_slice(&b).unwrap(); + assert_eq!(v["ok"], true); + assert_eq!(v["size"], 1 << 20); + assert_eq!(v["parts"], 1); + assert_eq!(v["enough"], true); + assert_eq!(v["dir"], models(&home).to_string_lossy().as_ref()); + + // Lancement, suivi, puis la sonde dit « déjà là ». + let (code, b) = fakego::send(facade, "POST", "/api/models/download", &[], &body).await; + assert_eq!(code, StatusCode::OK, "{b:?}"); + let mut done = false; + for _ in 0..500 { + let (_, _, b) = fakego::get(facade, "/api/models/download/status", &[]).await; + let v: Value = serde_json::from_slice(&b).unwrap(); + if v[0]["finished"] == true { + assert_eq!(v[0]["error"], ""); + assert_eq!(v[0]["done"], 1 << 20); + done = true; + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!(done); + let (_, b) = fakego::send(facade, "POST", "/api/models/download/probe", &[], &body).await; + let v: Value = serde_json::from_slice(&b).unwrap(); + assert_eq!( + (v["exists"].as_bool(), v["value"].as_str()), + (Some(true), Some("m.gguf")) + ); + assert_eq!(v["size"], 1 << 20); + + // Annuler l'inconnu ; corps vide. + let (code, b) = fakego::send( + facade, + "POST", + "/api/models/download/cancel", + &[], + r#"{"filename":"../x.gguf"}"#, + ) + .await; + assert_eq!(code, StatusCode::NOT_FOUND); + assert_eq!( + serde_json::from_slice::(&b).unwrap()["error"], + "aucun téléchargement pour x.gguf" + ); + let (code, b) = fakego::send(facade, "POST", "/api/models/download", &[], "").await; + assert_eq!( + ( + code, + serde_json::from_slice::(&b).unwrap()["error"].as_str() + ), + (StatusCode::BAD_REQUEST, Some("EOF")) + ); + + // Clé de pilotage posée : refus sans elle, la remise du Go passe. + go.set("state", "web_key", "cle"); + let (code, _, _) = fakego::get(facade, "/api/models/download/status", &[]).await; + assert_eq!(code, StatusCode::UNAUTHORIZED); + let (code, _, _) = fakego::get( + facade, + "/api/models/download/status", + &[("x-loki-internal", fakego::JETON)], + ) + .await; + assert_eq!(code, StatusCode::OK); + // Rien n'est parti au Go. + assert!( + !go.relayed + .lock() + .unwrap() + .iter() + .any(|p| p.contains("download")), + "{:?}", + go.relayed.lock().unwrap() + ); +} + +#[tokio::test] +async fn sans_la_variable_les_routes_restent_au_go() { + let tmp = tempfile::tempdir().unwrap(); + let (facade, go, _stop) = fakego::start(tmp.path().to_path_buf()).await; + let (_, _, b) = fakego::get(facade, "/api/models/download/status", &[]).await; + assert_eq!(&b[..], b"go:/api/models/download/status"); + assert!( + go.relayed + .lock() + .unwrap() + .iter() + .any(|p| p == "/api/models/download/status") + ); +} + +#[test] +fn part_perimes_vides_dans_chaque_dossier() { + let a = tempfile::tempdir().unwrap(); + let b = tempfile::tempdir().unwrap(); + for d in [a.path(), b.path()] { + std::fs::write(d.join("x.gguf.part"), b"x").unwrap(); + std::fs::write(d.join("X.GGUF.PART"), b"x").unwrap(); + std::fs::write(d.join("y.gguf"), b"y").unwrap(); + } + downloads::clean_stale_part_files(&[ + a.path().to_string_lossy().into_owned(), + b.path().to_string_lossy().into_owned(), + "/nexiste/pas".into(), + ]); + for d in [a.path(), b.path()] { + // Le motif du Go est sensible à la casse. + assert_eq!(entries(d), ["X.GGUF.PART", "y.gguf"]); + } +} diff --git a/rust/loki-rs/tests/lot2.rs b/rust/loki-rs/tests/lot2.rs index 5c789fa..0453b79 100644 --- a/rust/loki-rs/tests/lot2.rs +++ b/rust/loki-rs/tests/lot2.rs @@ -94,6 +94,7 @@ async fn banc(cle: Option<&str>) -> Banc { service: "loki-engine".into(), }, chat: None, + downloads: None, }; tokio::spawn(loki_rs::proxy::serve(l, cfg, async move { let _ = rx.await; diff --git a/rust/loki-rs/tests/parity_models/download.json b/rust/loki-rs/tests/parity_models/download.json new file mode 100644 index 0000000..ef9c430 --- /dev/null +++ b/rust/loki-rs/tests/parity_models/download.json @@ -0,0 +1,314 @@ +{ + "normalize": [ + { + "in": "", + "err": "lien vide" + }, + { + "in": " ", + "err": "lien vide" + }, + { + "in": "ftp://h/a.gguf", + "err": "lien invalide (http/https attendu)" + }, + { + "in": "pas un lien", + "err": "lien invalide (http/https attendu)" + }, + { + "in": "https://exemple.org/dossier/", + "err": "le lien doit pointer vers un fichier .gguf" + }, + { + "in": "https://huggingface.co/u/r/blob/main/m.gguf", + "url": "https://huggingface.co/u/r/resolve/main/m.gguf", + "name": "m.gguf", + "urls": [ + "https://huggingface.co/u/r/resolve/main/m.gguf" + ], + "names": [ + "m.gguf" + ] + }, + { + "in": "https://huggingface.co/u/r/blob/main/blob/m.gguf", + "url": "https://huggingface.co/u/r/resolve/main/blob/m.gguf", + "name": "m.gguf", + "urls": [ + "https://huggingface.co/u/r/resolve/main/blob/m.gguf" + ], + "names": [ + "m.gguf" + ] + }, + { + "in": "https://huggingface.co/u/r/resolve/main/m-Q4_K_M.gguf?download=true", + "url": "https://huggingface.co/u/r/resolve/main/m-Q4_K_M.gguf?download=true", + "name": "m-Q4_K_M.gguf", + "urls": [ + "https://huggingface.co/u/r/resolve/main/m-Q4_K_M.gguf?download=true" + ], + "names": [ + "m-Q4_K_M.gguf" + ] + }, + { + "in": " https://huggingface.co/u/r/blob/main/m.bin ", + "err": "le lien doit pointer vers un fichier .gguf" + }, + { + "in": "http://exemple.org/modeles/M.GGUF", + "url": "http://exemple.org/modeles/M.GGUF", + "name": "M.GGUF", + "urls": [ + "http://exemple.org/modeles/M.GGUF" + ], + "names": [ + "M.GGUF" + ] + }, + { + "in": "https://huggingface.co/u/r/blob/main/Q4/m-Q4-00002-of-00003.gguf", + "url": "https://huggingface.co/u/r/resolve/main/Q4/m-Q4-00002-of-00003.gguf", + "name": "m-Q4-00002-of-00003.gguf", + "urls": [ + "https://huggingface.co/u/r/resolve/main/Q4/m-Q4-00001-of-00003.gguf", + "https://huggingface.co/u/r/resolve/main/Q4/m-Q4-00002-of-00003.gguf", + "https://huggingface.co/u/r/resolve/main/Q4/m-Q4-00003-of-00003.gguf" + ], + "names": [ + "m-Q4-00001-of-00003.gguf", + "m-Q4-00002-of-00003.gguf", + "m-Q4-00003-of-00003.gguf" + ] + }, + { + "in": "https://huggingface.co/u/r/resolve/main/m-00001-of-00012.gguf", + "url": "https://huggingface.co/u/r/resolve/main/m-00001-of-00012.gguf", + "name": "m-00001-of-00012.gguf", + "urls": [ + "https://huggingface.co/u/r/resolve/main/m-00001-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00002-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00003-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00004-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00005-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00006-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00007-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00008-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00009-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00010-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00011-of-00012.gguf", + "https://huggingface.co/u/r/resolve/main/m-00012-of-00012.gguf" + ], + "names": [ + "m-00001-of-00012.gguf", + "m-00002-of-00012.gguf", + "m-00003-of-00012.gguf", + "m-00004-of-00012.gguf", + "m-00005-of-00012.gguf", + "m-00006-of-00012.gguf", + "m-00007-of-00012.gguf", + "m-00008-of-00012.gguf", + "m-00009-of-00012.gguf", + "m-00010-of-00012.gguf", + "m-00011-of-00012.gguf", + "m-00012-of-00012.gguf" + ] + }, + { + "in": "https://exemple.org/m-1-of-2.gguf", + "url": "https://exemple.org/m-1-of-2.gguf", + "name": "m-1-of-2.gguf", + "urls": [ + "https://exemple.org/m-1-of-2.gguf" + ], + "names": [ + "m-1-of-2.gguf" + ] + }, + { + "in": "https://exemple.org/m-02-of-01.gguf", + "url": "https://exemple.org/m-02-of-01.gguf", + "name": "m-02-of-01.gguf", + "urls": [ + "https://exemple.org/m-02-of-01.gguf" + ], + "names": [ + "m-02-of-01.gguf" + ] + } + ], + "content_range": { + "": 0, + "bytes 0-0/ 99 ": 99, + "bytes 0-0/*": 0, + "bytes 0-0/0": 0, + "bytes 0-0/12345": 12345 + }, + "source": [ + { + "status": 401, + "code": "GatedRepo", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "dépôt à accès restreint : accepte ses conditions sur huggingface.co/auteur/depot — puis renseigne le jeton Hugging Face (réglage « Jeton Hugging Face », ou variable d'environnement HF_TOKEN)" + }, + { + "status": 401, + "code": "GatedRepo", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": true, + "err": "dépôt à accès restreint : accepte ses conditions sur huggingface.co/auteur/depot — le jeton Hugging Face enregistré n'y donne pas accès (expiré, ou conditions non acceptées avec ce compte)" + }, + { + "status": 403, + "code": "RepoNotFound", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "dépôt privé ou inexistant : auteur/depot — puis renseigne le jeton Hugging Face (réglage « Jeton Hugging Face », ou variable d'environnement HF_TOKEN)" + }, + { + "status": 401, + "code": "GatedRepo", + "url": "https://exemple.org/m.gguf", + "token_set": false, + "err": "dépôt à accès restreint : ses conditions doivent être acceptées sur huggingface.co — puis renseigne le jeton Hugging Face (réglage « Jeton Hugging Face », ou variable d'environnement HF_TOKEN)" + }, + { + "status": 403, + "code": "", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": true, + "err": "accès refusé par Hugging Face sur auteur/depot — le jeton Hugging Face enregistré n'y donne pas accès (expiré, ou conditions non acceptées avec ce compte)" + }, + { + "status": 403, + "code": "", + "url": "https://exemple.org/m.gguf", + "token_set": false, + "err": "accès refusé par la source — puis renseigne le jeton Hugging Face (réglage « Jeton Hugging Face », ou variable d'environnement HF_TOKEN)" + }, + { + "status": 404, + "code": "EntryNotFound", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "fichier « m.gguf » absent du dépôt (HTTP 404) — dépôt réorganisé depuis la copie du lien, ou fragment jamais publié (table PLE livrée à part ?)" + }, + { + "status": 404, + "code": "EntryNotFound", + "url": "https://exemple.org/", + "token_set": false, + "err": "fichier absent du dépôt (HTTP 404) — dépôt réorganisé depuis la copie du lien, ou fragment jamais publié (table PLE livrée à part ?)" + }, + { + "status": 404, + "code": "RevisionNotFound", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "révision introuvable dans le dépôt (HTTP 404)" + }, + { + "status": 404, + "code": "", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "lien introuvable (HTTP 404)" + }, + { + "status": 416, + "code": "", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "la source refuse la plage d'octets demandée (HTTP 416) — fichier modifié pendant le transfert ?" + }, + { + "status": 429, + "code": "", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "trop de requêtes vers la source (HTTP 429) — réessaie dans quelques minutes" + }, + { + "status": 502, + "code": "", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "la source est en panne (HTTP 502) — réessaie plus tard" + }, + { + "status": 418, + "code": "", + "url": "https://huggingface.co/auteur/depot/resolve/main/m.gguf", + "token_set": false, + "err": "HTTP 418 depuis la source" + } + ], + "hf_host": { + "HUGGINGFACE.CO:443": true, + "cdn-lfs.huggingface.co": true, + "evilhuggingface.co": false, + "exemple.org": false, + "hf.co": true, + "huggingface.co": true, + "huggingface.co.evil.org": false, + "x.hf.co": true + }, + "hf_repo": { + "https://exemple.org/a/b/resolve/main/x.gguf": "", + "https://huggingface.co/-a/b/resolve/main/x.gguf": "", + "https://huggingface.co/a/b": "", + "https://huggingface.co/a/b/blob/main/x.gguf": "a/b", + "https://huggingface.co/a/b/tree/main": "", + "https://huggingface.co/auteur/depot/resolve/main/m.gguf": "auteur/depot" + }, + "dest": [ + { + "name": "a.gguf", + "dir": "", + "path": "{{HOME}}/models/a.gguf" + }, + { + "name": " ../../x.gguf ", + "dir": "", + "path": "{{HOME}}/models/x.gguf" + }, + { + "name": "x.bin", + "dir": "", + "err": "seuls les fichiers .gguf sont acceptés" + }, + { + "name": "", + "dir": "", + "err": "nom de modèle invalide" + }, + { + "name": "/", + "dir": "", + "err": "nom de modèle invalide" + }, + { + "name": "a.gguf", + "dir": "/etc", + "err": "dossier non autorisé : /etc — ajoute-le dans « Dossiers de modèles »" + }, + { + "name": "a.gguf", + "dir": "{{HOME}}/disque", + "path": "{{HOME}}/disque/a.gguf" + }, + { + "name": "a.gguf", + "dir": "{{HOME}}/disque/", + "path": "{{HOME}}/disque/a.gguf" + }, + { + "name": "a.gguf", + "dir": "{{HOME}}/models", + "path": "{{HOME}}/models/a.gguf" + } + ] +} diff --git a/rust/loki-rs/tests/proxy.rs b/rust/loki-rs/tests/proxy.rs index 575f6ca..2f5da3c 100644 --- a/rust/loki-rs/tests/proxy.rs +++ b/rust/loki-rs/tests/proxy.rs @@ -100,6 +100,7 @@ async fn facade(backend: SocketAddr) -> (SocketAddr, oneshot::Sender<()>) { service: "loki-engine".into(), }, chat: None, + downloads: None, }; tokio::spawn(loki_rs::proxy::serve(l, cfg, async move { let _ = rx.await;