From 76e65355b4c6476ebeb924c44d4d75cd273f212b Mon Sep 17 00:00:00 2001 From: Jared Ping <13906915+JaredP94@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:00:40 +0200 Subject: [PATCH] feat(proxies): route Anthropic Messages requests to model owners Port NVIDIA/Personal-AI-Router@7b21d1f from NVIDIA/Personal-AI-Router#27, originally addressing NVIDIA/Personal-AI-Router#16, to the fork's three engine-specific proxies. Ollama 0.14.0, LM Studio 0.4.1, and oMLX 0.1.0 all expose the Messages endpoint in their documented APIs. Classify POST /v1/messages as inference so it uses existing model eligibility, scheduling, failover, and first-body stall cooldown without translation. Keep oMLX's OpenAI usage-field injection off Messages requests so streamed Anthropic bodies pass through unchanged. Cover route roles, owner filtering, header and body passthrough, JSON and unbuffered SSE failover, stall cooldown, post-commit truncation, and older engine errors. Extend strict cross-process routing for all three engines, document compatibility, and bump each proxy to the additive 1.5.0 version. Co-authored-by: Ratherrayees <145657835+Ratherrayees@users.noreply.github.com> Signed-off-by: Jared Ping <13906915+JaredP94@users.noreply.github.com> --- docs/getting-started.mdx | 15 +- services/lmstudio-proxy/README.md | 21 ++ services/lmstudio-proxy/messages_test.go | 237 +++++++++++++++++++ services/lmstudio-proxy/proxy.go | 1 + services/ollama-proxy/README.md | 21 ++ services/ollama-proxy/messages_test.go | 237 +++++++++++++++++++ services/ollama-proxy/proxy.go | 1 + services/omlx-proxy/README.md | 21 ++ services/omlx-proxy/messages_test.go | 237 +++++++++++++++++++ services/omlx-proxy/proxy.go | 12 +- services/tests/model_routing_interop_test.go | 50 +++- services/versions.json | 6 +- 12 files changed, 846 insertions(+), 13 deletions(-) create mode 100644 services/lmstudio-proxy/messages_test.go create mode 100644 services/ollama-proxy/messages_test.go create mode 100644 services/omlx-proxy/messages_test.go diff --git a/docs/getting-started.mdx b/docs/getting-started.mdx index daa3ce7c7..71ec8cad8 100644 --- a/docs/getting-started.mdx +++ b/docs/getting-started.mdx @@ -461,8 +461,19 @@ What that base URL serves depends on the engine behind it: | Endpoint | Default base URL | Paths it serves | | --- | --- | --- | -| Ollama | `http://127.0.0.1:11434` | Ollama's own API — `/api/chat`, `/api/generate`, `/api/embed`, `/api/tags` — and the OpenAI-compatible `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/models` | -| LM Studio | `http://127.0.0.1:1234` | The OpenAI-compatible paths only | +| Ollama | `http://127.0.0.1:11434` | Ollama's own API — `/api/chat`, `/api/generate`, `/api/embed`, `/api/tags` — and the OpenAI-compatible `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/models`, plus Anthropic `POST /v1/messages` | +| LM Studio | `http://127.0.0.1:1234` | The OpenAI-compatible paths and Anthropic `POST /v1/messages` | + +Anthropic clients such as Claude Code can use `POST /v1/messages` with Ollama +0.14.0+, LM Studio 0.4.1+, or oMLX 0.1.0+. PAIR routes the body's `model` to an +advertised owner and passes JSON or streamed SSE through without translation. +Use the copied endpoint's API root (without a trailing `/v1`) as +`ANTHROPIC_BASE_URL`; the client appends `/v1/messages`. Older engines without +the route return an upstream HTTP error after eligible-owner failover is +exhausted; upgrade the serving engine. See the proxy references for +[Ollama](../services/ollama-proxy/README.md#anthropic-messages-api), +[LM Studio](../services/lmstudio-proxy/README.md#anthropic-messages-api), and +[oMLX](../services/omlx-proxy/README.md#anthropic-messages-api). The distinction matters when you fill in an application's settings. A client written against the OpenAI API usually wants the `/v1` included, as in diff --git a/services/lmstudio-proxy/README.md b/services/lmstudio-proxy/README.md index 160ee4227..8c3a85775 100644 --- a/services/lmstudio-proxy/README.md +++ b/services/lmstudio-proxy/README.md @@ -9,6 +9,27 @@ A discovery-aware HTTP reverse proxy for LM Studio nodes on the local network. I > **Clone of `ollama-proxy`.** This proxy is a deliberate clone of [`ollama-proxy`](../ollama-proxy/README.md) so the two share identical routing, failover, CORS, and node-selection behavior — the CORS policy is literally the same code, `nvpair-shared/cors`, and is documented [there](../ollama-proxy/README.md#http-reverse-proxy). The differences are engine-specific: it subscribes to the discovery relay for `lm` nodes, forwards the OpenAI-compatible inference routes (`/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`), tags workloads `lmstudio`, and persists its port to its own file. It has no `--alias-address`, so its self-forward guard covers only its own listener. +## Anthropic Messages API + +`POST /v1/messages` is an inference route, supported with **LM Studio +0.4.1 or newer** (see the engine's [release announcement](https://lmstudio.ai/blog/claudecode)). +The body’s `model` selects eligible owners through the same selection, +scheduler, failover, and workload path as other inference routes. PAIR performs +no protocol translation: the request body and response, including SSE with +`stream: true`, pass through unchanged. SSE is flushed as it arrives; +the existing first-body deadline and cooldown rules apply, and no attempt is +retried after response bytes reach the client. + +`x-api-key` and `anthropic-version` follow existing request-header forwarding; +`x-api-key` is not PAIR authentication. Existing proxy authentication and +loopback/cluster policies are unchanged. + +PAIR does not negotiate engine API versions. An older engine without this +endpoint returns its HTTP error (typically `404`); that status may try the next +eligible model owner, but only within the existing finite candidate list. +If none supports the endpoint, the final engine error is returned, not translated +or redirected to another model. Upgrade the serving engines to use this route. + ## Worker control health An id-bearing `worker/health` request (no params) returns `{"ok":true}` with diff --git a/services/lmstudio-proxy/messages_test.go b/services/lmstudio-proxy/messages_test.go new file mode 100644 index 000000000..56c6b3004 --- /dev/null +++ b/services/lmstudio-proxy/messages_test.go @@ -0,0 +1,237 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bufio" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestMessagesRouteRole(t *testing.T) { + for _, tc := range []struct { + method, path string + inference bool + }{ + {http.MethodPost, "/v1/messages", true}, + {http.MethodGet, "/v1/messages", false}, + {http.MethodPut, "/v1/messages", false}, + {http.MethodPost, "/v1/messages/count_tokens", false}, + } { + if got := isInferenceRequest(tc.method, tc.path); got != tc.inference { + t.Errorf("%s %s inference = %v, want %v", tc.method, tc.path, got, tc.inference) + } + } +} + +// Exercise real HTTP copying: the first SSE event must arrive while the engine +// is still waiting to finish, including after an advertised owner's 404. +func TestMessagesModelRoutingAndFailover(t *testing.T) { + for _, streaming := range []bool{false, true} { + t.Run(map[bool]string{false: "json", true: "sse"}[streaming], func(t *testing.T) { + body := `{"model":"model","max_tokens":7,"messages":[{"role":"user","content":"test"}],"stream":` + + map[bool]string{false: "false", true: "true"}[streaming] + `}` + var missingHits, goodHits, wrongHits atomic.Int32 + wrong := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + wrongHits.Add(1) + })) + defer wrong.Close() + check := func(r *http.Request) { + t.Helper() + data, err := io.ReadAll(r.Body) + if err != nil || string(data) != body || r.URL.Path != "/v1/messages" { + t.Error("engine did not receive the unchanged Messages request") + } + if r.Header.Get("x-api-key") != "synthetic-test-key" || + r.Header.Get("anthropic-version") != "2023-06-01" { + t.Error("Anthropic request headers were not passed through") + } + } + missing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + missingHits.Add(1) + check(r) + w.WriteHeader(http.StatusNotFound) + })) + defer missing.Close() + release := make(chan struct{}) + var once sync.Once + finish := func() { once.Do(func() { close(release) }) } + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + check(r) + if streaming { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_start\ndata: {}\n\n") + w.(http.Flusher).Flush() + <-release + _, _ = io.WriteString(w, "event: message_stop\ndata: {}\n\n") + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"type":"message","content":[]}`) + } + })) + defer good.Close() + defer finish() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "wrong", wrong.URL, "different-model")) + disc.AddManual(nodeForModel(t, "missing", missing.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + p.SetSelected("wrong") + p.SetPriority([]string{"missing", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer func() { + finish() + server.Close() + }() + client := &http.Client{Timeout: 3 * time.Second} + defer client.CloseIdleConnections() + req, err := http.NewRequest(http.MethodPost, server.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-api-key", "synthetic-test-key") + req.Header.Set("anthropic-version", "2023-06-01") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if streaming { + if resp.Header.Get("Content-Type") != "text/event-stream" { + t.Fatal("SSE content type changed") + } + reader := bufio.NewReader(resp.Body) + line, err := reader.ReadString('\n') + if err != nil || line != "event: message_start\n" { + t.Fatal("first SSE event did not arrive before engine completion") + } + finish() + rest, err := io.ReadAll(reader) + if err != nil || string(rest) != "data: {}\n\nevent: message_stop\ndata: {}\n\n" { + t.Fatal("SSE response changed") + } + } else { + data, err := io.ReadAll(resp.Body) + if err != nil || string(data) != `{"type":"message","content":[]}` { + t.Fatal("JSON response changed") + } + } + if missingHits.Load() != 1 || goodHits.Load() != 1 || wrongHits.Load() != 0 { + t.Fatal("request did not fail over only between advertised model owners") + } + }) + } +} + +func TestMessagesFirstBodyStallCooldown(t *testing.T) { + oldFirst, oldHeader := proxyFirstBodyTimeout, proxyResponseTimeout + proxyFirstBodyTimeout, proxyResponseTimeout = 50*time.Millisecond, 3*time.Second + defer func() { proxyFirstBodyTimeout, proxyResponseTimeout = oldFirst, oldHeader }() + var stalledHits atomic.Int32 + stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + stalledHits.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer stalled.Close() + const event = "event: message_stop\ndata: {}\n\n" + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, event) + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "stalled", stalled.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + for i := 0; i < 2; i++ { + p.SetPriority([]string{"stalled", "good"}) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"model","stream":true}`)).WithContext(ctx)) + cancel() + if rec.Code != http.StatusOK || rec.Body.String() != event { + t.Fatal("SSE first-body stall did not fail over") + } + if !p.coolingSet(time.Now())["stalled"] { + t.Fatal("first-body stall did not cool candidate") + } + } + if stalledHits.Load() != 1 { + t.Fatal("second request retried cooling candidate before a ready peer") + } +} + +func TestMessagesCommittedSSENeverRetries(t *testing.T) { + var goodHits atomic.Int32 + const partial = "event: message_start\ndata: {}\n\n" + broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Content-Length", "999") + _, _ = io.WriteString(w, partial) + w.(http.Flusher).Flush() + })) + defer broken.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { goodHits.Add(1) })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "broken", broken.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + p.SetPriority([]string{"broken", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer server.Close() + client := &http.Client{Timeout: 3 * time.Second} + defer client.CloseIdleConnections() + resp, err := client.Post(server.URL+"/v1/messages", "application/json", strings.NewReader(`{"model":"model","stream":true}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err == nil || string(data) != partial || goodHits.Load() != 0 { + t.Fatal("committed SSE was retried or not truncated") + } +} + +func TestMessagesOlderEngineReturnsError(t *testing.T) { + var hits atomic.Int32 + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusNotFound) + _, _ = io.WriteString(w, `{"error":"unsupported endpoint"}`) + })) + defer engine.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "old", engine.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"model","stream":true}`)).WithContext(ctx)) + if rec.Code != http.StatusNotFound || rec.Body.String() != `{"error":"unsupported endpoint"}` || hits.Load() != 1 || ctx.Err() != nil { + t.Fatal("older engine error was not returned cleanly after bounded candidates") + } +} diff --git a/services/lmstudio-proxy/proxy.go b/services/lmstudio-proxy/proxy.go index b41ddab99..af45470c0 100644 --- a/services/lmstudio-proxy/proxy.go +++ b/services/lmstudio-proxy/proxy.go @@ -171,6 +171,7 @@ var inferenceEndpoints = map[string]bool{ "/v1/chat/completions": true, "/v1/completions": true, "/v1/embeddings": true, + "/v1/messages": true, } // isInferenceRequest reports whether a request should be tracked as a diff --git a/services/ollama-proxy/README.md b/services/ollama-proxy/README.md index aa1c1e2f3..2fa1e197f 100644 --- a/services/ollama-proxy/README.md +++ b/services/ollama-proxy/README.md @@ -7,6 +7,27 @@ SPDX-License-Identifier: Apache-2.0 A discovery-aware HTTP reverse proxy for Ollama nodes on the local network. It runs no mDNS browse of its own: its routing targets come from the broker's discovery relay (it sends `discovery:subscribe {services:[ol]}` and replaces its routing overlay from each pushed `discovery:nodes` snapshot) plus user-added manual nodes. It forwards HTTP requests to the selected node, aggregates the model-list routes across candidate nodes, and exposes a bidirectional JSON-RPC 2.0 control channel over stdio (or an IPC socket). +## Anthropic Messages API + +`POST /v1/messages` is an inference route, supported with **Ollama +0.14.0 or newer** (see the engine's [route registration](https://github.com/ollama/ollama/blob/v0.14.0/server/routes.go)). +The body’s `model` selects eligible owners through the same selection, +scheduler, failover, and workload path as other inference routes. PAIR performs +no protocol translation: the request body and response, including SSE with +`stream: true`, pass through unchanged. SSE is flushed as it arrives; +the existing first-body deadline and cooldown rules apply, and no attempt is +retried after response bytes reach the client. + +`x-api-key` and `anthropic-version` follow existing request-header forwarding; +`x-api-key` is not PAIR authentication. Existing proxy authentication and +loopback/cluster policies are unchanged. + +PAIR does not negotiate engine API versions. An older engine without this +endpoint returns its HTTP error (typically `404`); that status may try the next +eligible model owner, but only within the existing finite candidate list. +If none supports the endpoint, the final engine error is returned, not translated +or redirected to another model. Upgrade the serving engines to use this route. + ## Worker control health An id-bearing `worker/health` request (no params) returns `{"ok":true}` with diff --git a/services/ollama-proxy/messages_test.go b/services/ollama-proxy/messages_test.go new file mode 100644 index 000000000..56c6b3004 --- /dev/null +++ b/services/ollama-proxy/messages_test.go @@ -0,0 +1,237 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bufio" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestMessagesRouteRole(t *testing.T) { + for _, tc := range []struct { + method, path string + inference bool + }{ + {http.MethodPost, "/v1/messages", true}, + {http.MethodGet, "/v1/messages", false}, + {http.MethodPut, "/v1/messages", false}, + {http.MethodPost, "/v1/messages/count_tokens", false}, + } { + if got := isInferenceRequest(tc.method, tc.path); got != tc.inference { + t.Errorf("%s %s inference = %v, want %v", tc.method, tc.path, got, tc.inference) + } + } +} + +// Exercise real HTTP copying: the first SSE event must arrive while the engine +// is still waiting to finish, including after an advertised owner's 404. +func TestMessagesModelRoutingAndFailover(t *testing.T) { + for _, streaming := range []bool{false, true} { + t.Run(map[bool]string{false: "json", true: "sse"}[streaming], func(t *testing.T) { + body := `{"model":"model","max_tokens":7,"messages":[{"role":"user","content":"test"}],"stream":` + + map[bool]string{false: "false", true: "true"}[streaming] + `}` + var missingHits, goodHits, wrongHits atomic.Int32 + wrong := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + wrongHits.Add(1) + })) + defer wrong.Close() + check := func(r *http.Request) { + t.Helper() + data, err := io.ReadAll(r.Body) + if err != nil || string(data) != body || r.URL.Path != "/v1/messages" { + t.Error("engine did not receive the unchanged Messages request") + } + if r.Header.Get("x-api-key") != "synthetic-test-key" || + r.Header.Get("anthropic-version") != "2023-06-01" { + t.Error("Anthropic request headers were not passed through") + } + } + missing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + missingHits.Add(1) + check(r) + w.WriteHeader(http.StatusNotFound) + })) + defer missing.Close() + release := make(chan struct{}) + var once sync.Once + finish := func() { once.Do(func() { close(release) }) } + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + check(r) + if streaming { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_start\ndata: {}\n\n") + w.(http.Flusher).Flush() + <-release + _, _ = io.WriteString(w, "event: message_stop\ndata: {}\n\n") + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"type":"message","content":[]}`) + } + })) + defer good.Close() + defer finish() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "wrong", wrong.URL, "different-model")) + disc.AddManual(nodeForModel(t, "missing", missing.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + p.SetSelected("wrong") + p.SetPriority([]string{"missing", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer func() { + finish() + server.Close() + }() + client := &http.Client{Timeout: 3 * time.Second} + defer client.CloseIdleConnections() + req, err := http.NewRequest(http.MethodPost, server.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-api-key", "synthetic-test-key") + req.Header.Set("anthropic-version", "2023-06-01") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if streaming { + if resp.Header.Get("Content-Type") != "text/event-stream" { + t.Fatal("SSE content type changed") + } + reader := bufio.NewReader(resp.Body) + line, err := reader.ReadString('\n') + if err != nil || line != "event: message_start\n" { + t.Fatal("first SSE event did not arrive before engine completion") + } + finish() + rest, err := io.ReadAll(reader) + if err != nil || string(rest) != "data: {}\n\nevent: message_stop\ndata: {}\n\n" { + t.Fatal("SSE response changed") + } + } else { + data, err := io.ReadAll(resp.Body) + if err != nil || string(data) != `{"type":"message","content":[]}` { + t.Fatal("JSON response changed") + } + } + if missingHits.Load() != 1 || goodHits.Load() != 1 || wrongHits.Load() != 0 { + t.Fatal("request did not fail over only between advertised model owners") + } + }) + } +} + +func TestMessagesFirstBodyStallCooldown(t *testing.T) { + oldFirst, oldHeader := proxyFirstBodyTimeout, proxyResponseTimeout + proxyFirstBodyTimeout, proxyResponseTimeout = 50*time.Millisecond, 3*time.Second + defer func() { proxyFirstBodyTimeout, proxyResponseTimeout = oldFirst, oldHeader }() + var stalledHits atomic.Int32 + stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + stalledHits.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer stalled.Close() + const event = "event: message_stop\ndata: {}\n\n" + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, event) + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "stalled", stalled.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + for i := 0; i < 2; i++ { + p.SetPriority([]string{"stalled", "good"}) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"model","stream":true}`)).WithContext(ctx)) + cancel() + if rec.Code != http.StatusOK || rec.Body.String() != event { + t.Fatal("SSE first-body stall did not fail over") + } + if !p.coolingSet(time.Now())["stalled"] { + t.Fatal("first-body stall did not cool candidate") + } + } + if stalledHits.Load() != 1 { + t.Fatal("second request retried cooling candidate before a ready peer") + } +} + +func TestMessagesCommittedSSENeverRetries(t *testing.T) { + var goodHits atomic.Int32 + const partial = "event: message_start\ndata: {}\n\n" + broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Content-Length", "999") + _, _ = io.WriteString(w, partial) + w.(http.Flusher).Flush() + })) + defer broken.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { goodHits.Add(1) })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "broken", broken.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + p.SetPriority([]string{"broken", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer server.Close() + client := &http.Client{Timeout: 3 * time.Second} + defer client.CloseIdleConnections() + resp, err := client.Post(server.URL+"/v1/messages", "application/json", strings.NewReader(`{"model":"model","stream":true}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err == nil || string(data) != partial || goodHits.Load() != 0 { + t.Fatal("committed SSE was retried or not truncated") + } +} + +func TestMessagesOlderEngineReturnsError(t *testing.T) { + var hits atomic.Int32 + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusNotFound) + _, _ = io.WriteString(w, `{"error":"unsupported endpoint"}`) + })) + defer engine.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "old", engine.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"model","stream":true}`)).WithContext(ctx)) + if rec.Code != http.StatusNotFound || rec.Body.String() != `{"error":"unsupported endpoint"}` || hits.Load() != 1 || ctx.Err() != nil { + t.Fatal("older engine error was not returned cleanly after bounded candidates") + } +} diff --git a/services/ollama-proxy/proxy.go b/services/ollama-proxy/proxy.go index 2773a9a9f..3b1b19a5a 100644 --- a/services/ollama-proxy/proxy.go +++ b/services/ollama-proxy/proxy.go @@ -177,6 +177,7 @@ var inferenceEndpoints = map[string]bool{ "/v1/chat/completions": true, "/v1/completions": true, "/v1/embeddings": true, + "/v1/messages": true, } // isInferenceRequest reports whether a request should be tracked as a diff --git a/services/omlx-proxy/README.md b/services/omlx-proxy/README.md index 409496503..74b182c2f 100644 --- a/services/omlx-proxy/README.md +++ b/services/omlx-proxy/README.md @@ -9,6 +9,27 @@ A discovery-aware HTTP reverse proxy for oMLX nodes on the local network (Apple > **OpenAI API Compatibility.** This proxy shares identical routing, failover, CORS, and node-selection behavior with [`lmstudio-proxy`](../lmstudio-proxy/README.md) and [`ollama-proxy`](../ollama-proxy/README.md). It subscribes to the discovery relay for `om` nodes, forwards the OpenAI-compatible inference routes (`/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`), tags workloads `omlx`, and persists its port to `omlx-proxy-port.json`. +## Anthropic Messages API + +`POST /v1/messages` is an inference route, supported with **oMLX +0.1.0 or newer** (see the engine's [route registration](https://github.com/jundot/omlx/blob/v0.1.0/omlx/server.py)). +The body’s `model` selects eligible owners through the same selection, +scheduler, failover, and workload path as other inference routes. PAIR performs +no protocol translation: the request body and response, including SSE with +`stream: true`, pass through unchanged. SSE is flushed as it arrives; +the existing first-body deadline and cooldown rules apply, and no attempt is +retried after response bytes reach the client. + +`x-api-key` and `anthropic-version` follow existing request-header forwarding; +`x-api-key` is not PAIR authentication. Existing proxy authentication and +loopback/cluster policies are unchanged. + +PAIR does not negotiate engine API versions. An older engine without this +endpoint returns its HTTP error (typically `404`); that status may try the next +eligible model owner, but only within the existing finite candidate list. +If none supports the endpoint, the final engine error is returned, not translated +or redirected to another model. Upgrade the serving engines to use this route. + ## Worker control health An id-bearing `worker/health` request (no params) returns `{"ok":true}` with diff --git a/services/omlx-proxy/messages_test.go b/services/omlx-proxy/messages_test.go new file mode 100644 index 000000000..56c6b3004 --- /dev/null +++ b/services/omlx-proxy/messages_test.go @@ -0,0 +1,237 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bufio" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestMessagesRouteRole(t *testing.T) { + for _, tc := range []struct { + method, path string + inference bool + }{ + {http.MethodPost, "/v1/messages", true}, + {http.MethodGet, "/v1/messages", false}, + {http.MethodPut, "/v1/messages", false}, + {http.MethodPost, "/v1/messages/count_tokens", false}, + } { + if got := isInferenceRequest(tc.method, tc.path); got != tc.inference { + t.Errorf("%s %s inference = %v, want %v", tc.method, tc.path, got, tc.inference) + } + } +} + +// Exercise real HTTP copying: the first SSE event must arrive while the engine +// is still waiting to finish, including after an advertised owner's 404. +func TestMessagesModelRoutingAndFailover(t *testing.T) { + for _, streaming := range []bool{false, true} { + t.Run(map[bool]string{false: "json", true: "sse"}[streaming], func(t *testing.T) { + body := `{"model":"model","max_tokens":7,"messages":[{"role":"user","content":"test"}],"stream":` + + map[bool]string{false: "false", true: "true"}[streaming] + `}` + var missingHits, goodHits, wrongHits atomic.Int32 + wrong := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + wrongHits.Add(1) + })) + defer wrong.Close() + check := func(r *http.Request) { + t.Helper() + data, err := io.ReadAll(r.Body) + if err != nil || string(data) != body || r.URL.Path != "/v1/messages" { + t.Error("engine did not receive the unchanged Messages request") + } + if r.Header.Get("x-api-key") != "synthetic-test-key" || + r.Header.Get("anthropic-version") != "2023-06-01" { + t.Error("Anthropic request headers were not passed through") + } + } + missing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + missingHits.Add(1) + check(r) + w.WriteHeader(http.StatusNotFound) + })) + defer missing.Close() + release := make(chan struct{}) + var once sync.Once + finish := func() { once.Do(func() { close(release) }) } + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + goodHits.Add(1) + check(r) + if streaming { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "event: message_start\ndata: {}\n\n") + w.(http.Flusher).Flush() + <-release + _, _ = io.WriteString(w, "event: message_stop\ndata: {}\n\n") + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"type":"message","content":[]}`) + } + })) + defer good.Close() + defer finish() + + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "wrong", wrong.URL, "different-model")) + disc.AddManual(nodeForModel(t, "missing", missing.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + p.SetSelected("wrong") + p.SetPriority([]string{"missing", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer func() { + finish() + server.Close() + }() + client := &http.Client{Timeout: 3 * time.Second} + defer client.CloseIdleConnections() + req, err := http.NewRequest(http.MethodPost, server.URL+"/v1/messages", strings.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-api-key", "synthetic-test-key") + req.Header.Set("anthropic-version", "2023-06-01") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if streaming { + if resp.Header.Get("Content-Type") != "text/event-stream" { + t.Fatal("SSE content type changed") + } + reader := bufio.NewReader(resp.Body) + line, err := reader.ReadString('\n') + if err != nil || line != "event: message_start\n" { + t.Fatal("first SSE event did not arrive before engine completion") + } + finish() + rest, err := io.ReadAll(reader) + if err != nil || string(rest) != "data: {}\n\nevent: message_stop\ndata: {}\n\n" { + t.Fatal("SSE response changed") + } + } else { + data, err := io.ReadAll(resp.Body) + if err != nil || string(data) != `{"type":"message","content":[]}` { + t.Fatal("JSON response changed") + } + } + if missingHits.Load() != 1 || goodHits.Load() != 1 || wrongHits.Load() != 0 { + t.Fatal("request did not fail over only between advertised model owners") + } + }) + } +} + +func TestMessagesFirstBodyStallCooldown(t *testing.T) { + oldFirst, oldHeader := proxyFirstBodyTimeout, proxyResponseTimeout + proxyFirstBodyTimeout, proxyResponseTimeout = 50*time.Millisecond, 3*time.Second + defer func() { proxyFirstBodyTimeout, proxyResponseTimeout = oldFirst, oldHeader }() + var stalledHits atomic.Int32 + stalled := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + stalledHits.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer stalled.Close() + const event = "event: message_stop\ndata: {}\n\n" + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, event) + })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "stalled", stalled.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + for i := 0; i < 2; i++ { + p.SetPriority([]string{"stalled", "good"}) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"model","stream":true}`)).WithContext(ctx)) + cancel() + if rec.Code != http.StatusOK || rec.Body.String() != event { + t.Fatal("SSE first-body stall did not fail over") + } + if !p.coolingSet(time.Now())["stalled"] { + t.Fatal("first-body stall did not cool candidate") + } + } + if stalledHits.Load() != 1 { + t.Fatal("second request retried cooling candidate before a ready peer") + } +} + +func TestMessagesCommittedSSENeverRetries(t *testing.T) { + var goodHits atomic.Int32 + const partial = "event: message_start\ndata: {}\n\n" + broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Content-Length", "999") + _, _ = io.WriteString(w, partial) + w.(http.Flusher).Flush() + })) + defer broken.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { goodHits.Add(1) })) + defer good.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "broken", broken.URL, "model")) + disc.AddManual(nodeForModel(t, "good", good.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + p.SetPriority([]string{"broken", "good"}) + server := httptest.NewServer(http.HandlerFunc(p.handleHTTP)) + defer server.Close() + client := &http.Client{Timeout: 3 * time.Second} + defer client.CloseIdleConnections() + resp, err := client.Post(server.URL+"/v1/messages", "application/json", strings.NewReader(`{"model":"model","stream":true}`)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err == nil || string(data) != partial || goodHits.Load() != 0 { + t.Fatal("committed SSE was retried or not truncated") + } +} + +func TestMessagesOlderEngineReturnsError(t *testing.T) { + var hits atomic.Int32 + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusNotFound) + _, _ = io.WriteString(w, `{"error":"unsupported endpoint"}`) + })) + defer engine.Close() + disc := NewDiscovery() + disc.AddManual(nodeForModel(t, "old", engine.URL, "model")) + p := testProxy(disc, 0) + defer p.closeIdleTransports() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + rec := httptest.NewRecorder() + p.handleHTTP(rec, httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"model","stream":true}`)).WithContext(ctx)) + if rec.Code != http.StatusNotFound || rec.Body.String() != `{"error":"unsupported endpoint"}` || hits.Load() != 1 || ctx.Err() != nil { + t.Fatal("older engine error was not returned cleanly after bounded candidates") + } +} diff --git a/services/omlx-proxy/proxy.go b/services/omlx-proxy/proxy.go index 52de603cf..e9a89e6ab 100644 --- a/services/omlx-proxy/proxy.go +++ b/services/omlx-proxy/proxy.go @@ -173,6 +173,7 @@ var inferenceEndpoints = map[string]bool{ "/v1/chat/completions": true, "/v1/completions": true, "/v1/embeddings": true, + "/v1/messages": true, } // isInferenceRequest reports whether a request should be tracked as a @@ -1124,13 +1125,12 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { bodyBytes, model := bufferBodyAndModel(r) isInf := isInferenceRequest(r.Method, r.URL.Path) - // The body forwarded to the engine asks oMLX for its usage block, so the - // engine's prefill/decode numbers reach the telemetry reader (oMLX emits - // that block only when the request asks for it). The client's own body is - // never modified: the flag is added to the buffered copy that every - // failover attempt replays, so all attempts forward the same body. + // OpenAI inference asks oMLX for its usage block so prefill/decode numbers + // reach the telemetry reader. The flag is added to the buffered copy every + // failover attempt replays, never the client's body. Anthropic Messages + // requests pass through unchanged and must not get OpenAI stream_options. forwardBody := bodyBytes - if isInf { + if isInf && r.URL.Path != "/v1/messages" { if injected := withUsageRequested(bodyBytes); !bytes.Equal(injected, bodyBytes) { forwardBody = injected // The forwarded body is longer than the client's, so a declared diff --git a/services/tests/model_routing_interop_test.go b/services/tests/model_routing_interop_test.go index b6c13dd77..5255fb089 100644 --- a/services/tests/model_routing_interop_test.go +++ b/services/tests/model_routing_interop_test.go @@ -5,6 +5,7 @@ package tests import ( "bytes" + "encoding/json" "fmt" "io" "net/http" @@ -13,6 +14,8 @@ import ( "sync/atomic" "testing" "time" + + "nvpair-shared/jsonrpc" ) type routingUpstream struct { @@ -38,6 +41,44 @@ func newRoutingUpstream(t *testing.T, status int) *routingUpstream { return upstream } +// Match only the requested status response; startup notifications and earlier +// workers' replies must not supply another proxy's port. +func waitRoutingProxyReady(t *testing.T, stdin io.Writer, msgs <-chan jsonrpc.Message, prefix string) int { + t.Helper() + id := map[string]int{"proxy": 10000, "lmstudio-proxy": 20000, "omlx-proxy": 30000}[prefix] + send := func() { + writeRawFrame(t, stdin, fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":%q}`, id, prefix+":get-status")) + } + send() + deadline := time.NewTimer(15 * time.Second) + defer deadline.Stop() + tick := time.NewTicker(100 * time.Millisecond) + defer tick.Stop() + for { + select { + case msg, ok := <-msgs: + if !ok { + t.Fatal("broker stream closed waiting for proxy readiness") + } + if msg.ID == nil || string(*msg.ID) != fmt.Sprint(id) { + continue + } + var status struct { + Ready bool `json:"ready"` + Port int `json:"port"` + } + if json.Unmarshal(msg.Result, &status) == nil && status.Ready { + return status.Port + } + case <-tick.C: + id++ + send() + case <-deadline.C: + t.Fatalf("timed out waiting for %s readiness", prefix) + } + } +} + func TestStrictModelRoutingAcrossProcesses(t *testing.T) { if portBusy(11435) || portBusy(1234) { t.Skip("ollama-proxy (11435) or lmstudio-proxy (1234) default port already in use; skipping") @@ -50,6 +91,7 @@ func TestStrictModelRoutingAcrossProcesses(t *testing.T) { stdin, msgs, stderr, cleanup := startBrokerWith(t, "--proxy-path", proxyBin, "--lmstudio-proxy-path", lmstudioProxyBin, + "--omlx-proxy-path", omlxProxyBin, ) t.Cleanup(cleanup) go func() { @@ -58,8 +100,9 @@ func TestStrictModelRoutingAcrossProcesses(t *testing.T) { }() waitForMethod(t, msgs, "app:ready", 10*time.Second) - ollamaPort := waitProxyReady(t, stdin, msgs, 15*time.Second) - lmstudioPort := waitLMStudioProxyReady(t, stdin, msgs, 15*time.Second) + ollamaPort := waitRoutingProxyReady(t, stdin, msgs, "proxy") + lmstudioPort := waitRoutingProxyReady(t, stdin, msgs, "lmstudio-proxy") + omlxPort := waitRoutingProxyReady(t, stdin, msgs, "omlx-proxy") type proxyCase struct { name string @@ -69,7 +112,10 @@ func TestStrictModelRoutingAcrossProcesses(t *testing.T) { } cases := []proxyCase{ {name: "ollama", rpcPrefix: "proxy", path: "/api/chat", port: ollamaPort}, + {name: "ollama-anthropic", rpcPrefix: "proxy", path: "/v1/messages", port: ollamaPort}, {name: "lmstudio", rpcPrefix: "lmstudio-proxy", path: "/v1/chat/completions", port: lmstudioPort}, + {name: "lmstudio-anthropic", rpcPrefix: "lmstudio-proxy", path: "/v1/messages", port: lmstudioPort}, + {name: "omlx-anthropic", rpcPrefix: "omlx-proxy", path: "/v1/messages", port: omlxPort}, } client := &http.Client{Timeout: 5 * time.Second} t.Cleanup(client.CloseIdleConnections) diff --git a/services/versions.json b/services/versions.json index 0cc2e7abd..313bd59cf 100644 --- a/services/versions.json +++ b/services/versions.json @@ -3,8 +3,8 @@ "product": "2.6.0", "installer": "2.6.0", "components": { - "ollama-proxy": "1.4.1", - "lmstudio-proxy": "1.4.1", + "ollama-proxy": "1.5.0", + "lmstudio-proxy": "1.5.0", "nvpair-node-info": "1.4.0", "nvpair-node-scanner": "1.3.0", "nvpair-manual-nodes": "1.3.0", @@ -16,6 +16,6 @@ "nvpair-cluster-manager": "2.2.0", "nvpair-job-scheduler": "1.1.0", "nvpair-tui": "1.3.0", - "omlx-proxy": "1.4.2" + "omlx-proxy": "1.5.0" } }