Compare commits

..
Author SHA1 Message Date
Ilker ManapandClaude Opus 5.5 d1e91868ee Setup asks where alerts go, and sends each a test
notify.webhooks existed and setup never asked for one, so a new site
learnt of a parked model or a lost node from its users. `llmgw server
init --webhook URL` (repeatable) and a Notification webhooks field on
both setup screens take them now. Slack's address, Mattermost's and
Rocket.Chat's /hooks/ and Discord's …/slack are sent Slack's text, any
other the event as JSON; slack: or json: in front says which.

Once the control plane runs, they are added to the cluster's settings
beside the webhooks already there, none twice, and setup waits for the
control planes to start again on the change before sending each a test,
saying which answered. A webhook's path is its secret, so it is shown by
host only. A failure is said, not fatal: the cluster serves without its
alerts, and ctl notify test tries again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-14 13:55:48 +02:00
18 changed files with 453 additions and 1326 deletions

No files matched your search

+6 -11
View File
@@ -28,17 +28,12 @@ APIs change between minors without a deprecation period.
### Added
- **Pictures, 3D models, vision and workflows as MCP tools** at `/mcp`,
beside the memory notes, each offered only where the gateway has a model
for it: `generate_image` (with `concept`, `side` and `texture` styles),
`generate_mesh` (from a picture or a prompt), `review_images` (a vision
model, to a JSON schema), `run_workflow`, `job_status` and
`gateway_status`. Each makes its request under the caller's key, so
allowlists, limits, licences and budgets apply and it is charged like
any request; what it makes is kept as the caller's file and answered by
id and download URL. A client that takes a stream gets
`notifications/progress` while a mesh is made and waits up to ten
minutes; one that does not, 50 seconds, then the job id.
- **Setup asks where alerts go**: `llmgw server init --webhook URL`
(repeatable) and a **Notification webhooks** field on both setup screens.
Slack's, Mattermost's and Rocket.Chat's addresses are sent Slack's text,
any other the event as JSON (`slack:` or `json:` in front says which).
They are added to the cluster's `notify.webhooks` beside those already
there, and each is sent a test once the control plane runs on them.
- **Extended thinking on `/v1/messages`**: `thinking: {type: "enabled"}`
tells a model that thinks by its chat template to (`enable_thinking`),
+1 -1
View File
@@ -82,7 +82,7 @@
"/mcp": {
"post": {
"operationId": "postMcp",
"summary": "The gateway's MCP server: the memory notes, and pictures, 3D models, vision and workflows, as tools",
"summary": "The gateway's MCP server: the memory notes as tools",
"tags": [
"mcp"
],
-36
View File
@@ -534,42 +534,6 @@ curl https://llm.example.com/mcp \
-d '{"jsonrpc": "2.0", "id": 1, "method": "tools/list"}'
```
**Pictures, 3D models, vision and workflows.** Where the gateway has the
models for them, the same server offers more tools — only those it has a
model for:
| Tool | Does |
|---|---|
| `generate_image` | A picture from a prompt. `style` adds what a picture for a mesh wants: `concept` (one object, three-quarter view, grey background), `side` (an exact side view), `texture` (a seamless tile) or `none`. Answers the PNG's `file_id`, `url` and `seed`, with a 512 px preview to look at. |
| `generate_mesh` | A GLB from a picture — `image` is a `file-…` id (a `generate_image` result, say), a data URL or base64 — or from a `prompt`, for which a concept picture is made first. `seed`, `resolution`, `texture_resolution`, and any other input the model's workflow declares in `inputs`. Answers the GLB's `file_id`, `url` and licence. |
| `review_images` | Asks a vision model about one to four pictures (each shrunk to 1280 px on its long side); `json_schema` makes it answer that JSON. |
| `run_workflow` | Runs a ComfyUI workflow from the inputs its catalog entry declares, and answers its result file. |
| `job_status` | Where a mesh or a workflow run stands, and its file when done; `wait_seconds` waits for it. |
| `gateway_status` | The picture, 3D, vision and workflow models, whether each is loaded on a GPU now, a workflow's inputs, and your last jobs. |
Each tool makes its request of the gateway under your key, so your
allowlist, limits, left-out licences and budget apply, and it is charged
and listed like any other request. What it makes is kept as your file for
as long as job results are: fetch it from `url` with the same key (`curl -H
"Authorization: Bearer $LLMGW_API_KEY" -o rock.glb <url>`). A tool server
over HTTP cannot write to your disk, so an agent that needs the file
locally — to import it in Blender — fetches it, or goes through a small
local bridge that does.
A mesh takes from half a minute to a few minutes. A client that takes the
answer as a stream (`Accept: text/event-stream`, as Claude Code does) is
answered that way: `notifications/progress` for the call's
`progressToken` while it runs — queued, waiting for the GPU, the node and
its percent — then the result, and the call waits up to ten minutes. One
answered in one piece waits 50 seconds; a job not done by then is answered
with its `job_id`, to ask after with `job_status`. `wait_seconds` (up to
900) sets either.
```
claude mcp add --transport http llmgw https://llm.example.com/mcp \
--header "Authorization: Bearer $LLMGW_API_KEY"
```
### Licences you leave out
A model's licence is in `/v1/models` (`license`: its `id`, `name`, whether
+17
View File
@@ -1827,6 +1827,23 @@ same event about the same thing is sent at most once an hour, so a budget that
refuses a thousand requests is one message. `llmgw ctl notify test` sends a
test event to each webhook and says how each answered.
Setup asks for them: **Notification webhooks** on its first page, or
`--webhook` to `llmgw server init`, as often as there are webhooks:
```
sudo llmgw server init --admin ops@example.com \
--webhook https://hooks.slack.com/services/… \
--webhook json:https://ops.example.com/llmgw
```
An address of Slack's, or one with `/hooks/` in it (Mattermost, Rocket.Chat)
or ending in `/slack`, is sent Slack's `{"text": …}`; any other the event as
JSON; `slack:` or `json:` in front says which. They are added to the
cluster's settings beside any already there — run again, nothing is added
twice — and once the control plane runs on them, each is sent a test, whose
answer setup prints. Events to choose and a signing secret are set in the
settings afterwards (`llmgw ctl config edit`).
### Tracing
Metrics say that a model took four minutes to start; a trace says where they
+9
View File
@@ -73,6 +73,9 @@ type serverInit struct {
// patroni are user@host of the machines to keep the store on under
// Patroni, replicated; this machine is the third etcd vote for two.
patroni []string
// webhooks are where the cluster's alerts go: put into its settings
// once it runs, and each sent a test.
webhooks []string
// patroniNodes are the same machines already connected to, which the
// setup screen does with the password it was given; when set, patroni
// only names them.
@@ -182,6 +185,9 @@ func newServerInitCmd() *cobra.Command {
"leave the firewall alone, instead of keeping llmgw's internal ports to the cluster's networks")
f.BoolVar(&in.postgres, "postgres", false,
"keep the store in PostgreSQL on this machine, installing it and moving an existing SQLite store into it")
f.StringArrayVar(&in.webhooks, "webhook", nil,
"where alerts go — a model parked, a node lost, a budget spent: an http(s) address, Slack's, Mattermost's "+
"and Rocket.Chat's sent Slack's text, any other the event as JSON (slack: or json: in front says which); repeatable")
f.StringArrayVar(&in.patroni, "patroni", nil,
"user@host of a machine to keep the store on, replicated under Patroni; give two or more. "+
"Installs PostgreSQL, Patroni and etcd there over ssh, and moves the existing store in")
@@ -249,6 +255,9 @@ func (in *serverInit) run(ctx context.Context) error {
fmt.Fprintf(in.out, " !! firewall: %v; llmgw's ports answer every network this machine is on\n", err)
}
}
if err := in.ensureWebhooks(ctx); err != nil {
fmt.Fprintf(in.out, " !! webhooks: %v\n", err)
}
if in.judgeWanted() && in.configJudgesWith(guardModelName) {
// Not fatal: without its judge the baseline policy still checks for
// secrets and injection, and a judged request is served flagged.
+185
View File
@@ -0,0 +1,185 @@
package cli
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/url"
"slices"
"strings"
"time"
"go.yaml.in/yaml/v3"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
)
// Where the cluster says what went wrong: a model parked, a node lost, a
// budget spent. notify.webhooks existed, and setup never asked for one, so a
// site learnt of a lost node from its users.
// parseWebhook reads a webhook as setup takes it: an address, with "slack:"
// or "json:" in front to say what it is sent, else what its address says.
func parseWebhook(s string) (notify.Webhook, error) {
s = strings.TrimSpace(s)
format := ""
for _, f := range []string{"slack", "json"} {
if rest, ok := strings.CutPrefix(s, f+":"); ok && !strings.HasPrefix(rest, "//") {
format, s = f, rest
}
}
u, err := url.Parse(s)
if err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Host == "" {
return notify.Webhook{}, fmt.Errorf("webhook %q: an http(s) address, with slack: or json: in front to say its format", s)
}
if format == "" {
format = webhookFormat(u)
}
return notify.Webhook{URL: s, Format: format}, nil
}
// webhookFormat is what an address takes: Slack's own, and the incoming
// webhooks that read Slack's {"text": …} — Mattermost's and Rocket.Chat's
// /hooks/, Discord's …/slack — are sent that; anything else the event as
// JSON.
func webhookFormat(u *url.URL) string {
switch {
case u.Host == "hooks.slack.com",
strings.Contains(u.Path, "/hooks/"),
strings.HasSuffix(u.Path, "/slack"):
return "slack"
}
return "json"
}
// parseWebhooks reads a list of them, by commas or as given.
func parseWebhooks(list []string) ([]notify.Webhook, error) {
var out []notify.Webhook
for _, item := range list {
for _, s := range strings.Split(item, ",") {
if strings.TrimSpace(s) == "" {
continue
}
w, err := parseWebhook(s)
if err != nil {
return nil, err
}
out = append(out, w)
}
}
return out, nil
}
// webhookEntry is one of notify.webhooks as the settings document has it.
type webhookEntry struct {
URL string `yaml:"url"`
Format string `yaml:"format,omitempty"`
Events []string `yaml:"events,omitempty,flow"`
SecretFile string `yaml:"secret_file,omitempty"`
}
// withWebhooks adds these webhooks to a settings document's notify.webhooks,
// and says whether that changed it: one already there, by address, is left
// as it is, with whatever events and secret it was given.
func withWebhooks(doc *settings.Doc, add []notify.Webhook) (bool, error) {
path := []string{"notify", "webhooks"}
var list []webhookEntry
if n := doc.Node(path); n != nil {
if err := n.Decode(&list); err != nil {
return false, fmt.Errorf("notify.webhooks in the cluster's settings: %w", err)
}
}
changed := false
for _, w := range add {
if !slices.ContainsFunc(list, func(e webhookEntry) bool { return e.URL == w.URL }) {
list = append(list, webhookEntry{URL: w.URL, Format: w.Format})
changed = true
}
}
if !changed {
return false, nil
}
b, err := yaml.Marshal(list)
if err != nil {
return false, err
}
return true, doc.SetYAML(path, string(b))
}
// ensureWebhooks puts the webhooks given into the cluster's settings and
// sends each a test, once the control plane runs on them. Said, not fatal: a
// cluster without its alerts still serves.
func (in *serverInit) ensureWebhooks(ctx context.Context) error {
if len(in.webhooks) == 0 {
return nil
}
hooks, err := parseWebhooks(in.webhooks)
if err != nil {
return err
}
c := in.adminAPI()
var cur adminapi.SettingsView
if err := c.get(ctx, "/api/v1/settings", &cur); err != nil {
return err
}
doc, err := settings.FromBytes("settings", []byte(cur.Body))
if err != nil {
return err
}
changed, err := withWebhooks(doc, hooks)
if err != nil {
return err
}
if changed {
b, err := doc.Bytes()
if err != nil {
return err
}
req, _ := json.Marshal(adminapi.SettingsRequest{Base: cur.Version, Body: string(b), Note: "notification webhooks, from setup"})
if err := c.put(ctx, "/api/v1/settings", req, &cur); err != nil {
return err
}
in.done("notification webhooks in the cluster's settings (version %d)", cur.Version)
} else {
in.kept("notification webhooks (already in the cluster's settings)")
}
// The control planes start again on the change; the test is sent by
// one that knows every webhook asked for.
var res []notify.Result
err = waitFor(ctx, 2*time.Minute, "the control plane to run on its webhooks", func() (bool, error) {
res = nil
if c.post(ctx, "/api/v1/notify/test", nil, &res) != nil {
return false, nil
}
for _, w := range hooks {
// Answered by host: the path of a webhook is its secret.
if !slices.ContainsFunc(res, func(r notify.Result) bool { return r.URL == notify.Redact(w.URL) }) {
return false, nil
}
}
return true, nil
})
if err != nil {
return err
}
var failed []string
for _, r := range res {
if !slices.ContainsFunc(hooks, func(w notify.Webhook) bool { return notify.Redact(w.URL) == r.URL }) {
continue
}
if r.Err != "" {
failed = append(failed, r.URL+": "+r.Err)
continue
}
in.done("a test notification reached %s", r.URL)
}
if len(failed) > 0 {
return errors.New("a test notification did not arrive: " + strings.Join(failed, "; ") +
" (kept in the settings: 'llmgw ctl notify test' tries again)")
}
return nil
}
+122
View File
@@ -0,0 +1,122 @@
package cli
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
)
func TestAWebhooksFormatIsWhatItsAddressTakes(t *testing.T) {
for in, want := range map[string]string{
"https://hooks.slack.com/services/T1/B2/abc": "slack",
"https://chat.example.com/hooks/xyz": "slack", // Mattermost, Rocket.Chat
"https://discord.com/api/webhooks/1/abc/slack": "slack",
"https://alerts.example.com/llmgw": "json",
"json:https://hooks.slack.com/services/T1/B2/abc": "json",
"slack:https://alerts.example.com/llmgw": "slack",
" https://events.pagerduty.com/integration/x/enqueue ": "json",
} {
w, err := parseWebhook(in)
if err != nil || w.Format != want || strings.HasPrefix(w.URL, "slack:") || strings.HasPrefix(w.URL, "json:") {
t.Errorf("%q: %+v %v, want %s", in, w, err, want)
}
}
for _, bad := range []string{"hooks.slack.com/x", "ftp://example.com/x", "slack:", "https://"} {
if _, err := parseWebhook(bad); err == nil {
t.Errorf("%q was taken", bad)
}
}
}
func TestWebhooksAreAddedToTheSettingsBesideTheOnesThere(t *testing.T) {
doc, err := settings.FromBytes("settings", []byte("router:\n retries: 2\nnotify:\n webhooks:\n"+
" - url: https://alerts.example.com/a\n events: [node_lost]\n secret_file: /etc/llmgw/secrets/hook\n"))
if err != nil {
t.Fatal(err)
}
hooks, _ := parseWebhooks([]string{"https://alerts.example.com/a, https://hooks.slack.com/services/T/B/c"})
changed, err := withWebhooks(doc, hooks)
if err != nil || !changed {
t.Fatalf("changed %v: %v", changed, err)
}
b, _ := doc.Bytes()
got := string(b)
for _, want := range []string{"retries: 2", "url: https://alerts.example.com/a", "events: [node_lost]",
"secret_file: /etc/llmgw/secrets/hook", "url: https://hooks.slack.com/services/T/B/c", "format: slack"} {
if !strings.Contains(got, want) {
t.Errorf("%q not in:\n%s", want, got)
}
}
if strings.Count(got, "alerts.example.com/a") != 1 {
t.Errorf("a webhook already there was added again:\n%s", got)
}
if changed, _ := withWebhooks(doc, hooks); changed {
t.Error("adding the same webhooks again changed the settings")
}
}
func TestSetupPutsItsWebhooksInTheSettingsAndTestsThemOnceInForce(t *testing.T) {
var mu sync.Mutex
body, tests := "router:\n retries: 2\n", 0
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mu.Lock()
defer mu.Unlock()
w.Header().Set("Content-Type", "application/json")
switch {
case r.Method == http.MethodGet && r.URL.Path == "/api/v1/settings":
_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 1, Body: body})
case r.Method == http.MethodPut && r.URL.Path == "/api/v1/settings":
var req adminapi.SettingsRequest
_ = json.NewDecoder(r.Body).Decode(&req)
body = req.Body
_ = json.NewEncoder(w).Encode(adminapi.SettingsView{Version: 2, Body: body})
case r.Method == http.MethodPost && r.URL.Path == "/api/v1/notify/test":
tests++
// The first answer comes from the control plane before it
// started again on the change: it knows no webhooks.
res := []notify.Result{}
if tests > 1 {
res = []notify.Result{{URL: notify.Redact("https://hooks.slack.com/services/T/B/c")},
{URL: notify.Redact("https://ops.example.com/in"), Err: "502 Bad Gateway"}}
}
_ = json.NewEncoder(w).Encode(res)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "server.yaml"),
[]byte("admin_listen: "+strings.TrimPrefix(srv.URL, "http://")+"\n"), 0o600); err != nil {
t.Fatal(err)
}
var out bytes.Buffer
in := &serverInit{etcDir: dir, out: &out,
webhooks: []string{"https://hooks.slack.com/services/T/B/c", "https://ops.example.com/in"}}
err := in.ensureWebhooks(t.Context())
if err == nil || !strings.Contains(err.Error(), "502") || !strings.Contains(err.Error(), "ctl notify test") {
t.Errorf("a webhook that did not answer: %v", err)
}
mu.Lock()
defer mu.Unlock()
if !strings.Contains(body, "format: slack") || !strings.Contains(body, "url: https://ops.example.com/in") ||
!strings.Contains(body, "retries: 2") || tests < 2 {
t.Errorf("settings %q after %d tests", body, tests)
}
if strings.Contains(out.String(), "/services/T/B/c") {
t.Errorf("a webhook's secret path was printed:\n%s", out.String())
}
if !strings.Contains(out.String(), "reached https://hooks.slack.com/…") {
t.Errorf("the one that answered is not said:\n%s", out.String())
}
}
+28 -1
View File
@@ -50,6 +50,7 @@ func (w *lineWriter) Write(p []byte) (int, error) {
const (
fAdmin = iota
fHosts
fWebhooks
fPostgres
fMirror
fNode
@@ -70,6 +71,7 @@ type initModel struct {
focus int
admin textField
hosts textField
webhooks textField
bundle textField
bundleSHA textField
toggles [fCount]bool
@@ -103,6 +105,11 @@ func newInitModel(in *serverInit, found []string, askForm bool) initModel {
hosts.SetValue(strings.Join(in.hosts, ", "))
hosts.CharLimit = 400
hosts.Width = 40
webhooks := newTextField()
webhooks.Placeholder = "https://hooks.slack.com/services/…"
webhooks.SetValue(strings.Join(in.webhooks, ", "))
webhooks.CharLimit = 800
webhooks.Width = 40
bundle := newTextField()
bundle.Placeholder = "/media/usb/llmgw-bundle.tar"
bundle.SetValue(in.bundle)
@@ -114,7 +121,7 @@ func newInitModel(in *serverInit, found []string, askForm bool) initModel {
bundleSHA.CharLimit = 64
bundleSHA.Width = 64
sp := newSpinner()
m := initModel{in: in, found: found, form: askForm, admin: admin, hosts: hosts,
m := initModel{in: in, found: found, form: askForm, admin: admin, hosts: hosts, webhooks: webhooks,
bundle: bundle, bundleSHA: bundleSHA, spin: sp, width: 100}
m.toggles[fPostgres] = in.postgres
m.toggles[fMirror] = !in.noMirror
@@ -235,6 +242,17 @@ func (m initModel) formKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
m.moveFocus(0)
return m, nil
}
hooks := strings.TrimSpace(m.webhooks.Value())
if _, err := parseWebhooks([]string{hooks}); err != nil {
m.formError = err.Error()
m.focus = fWebhooks
m.moveFocus(0)
return m, nil
}
m.in.webhooks = nil
if hooks != "" {
m.in.webhooks = []string{hooks}
}
m.in.bundle, m.in.bundleSHA = bundle, sum
m.in.admin = admin
m.in.hosts = nil
@@ -260,6 +278,10 @@ func (m initModel) formKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
var cmd tea.Cmd
m.hosts, cmd = m.hosts.Update(k)
return m, cmd
case fWebhooks:
var cmd tea.Cmd
m.webhooks, cmd = m.webhooks.Update(k)
return m, cmd
case fBundle:
var cmd tea.Cmd
m.bundle, cmd = m.bundle.Update(k)
@@ -284,6 +306,7 @@ func (m *initModel) moveFocus(d int) {
}
m.admin.Blur()
m.hosts.Blur()
m.webhooks.Blur()
m.bundle.Blur()
m.bundleSHA.Blur()
switch m.focus {
@@ -291,6 +314,8 @@ func (m *initModel) moveFocus(d int) {
m.admin.Focus()
case fHosts:
m.hosts.Focus()
case fWebhooks:
m.webhooks.Focus()
case fBundle:
m.bundle.Focus()
case fBundleSHA:
@@ -409,6 +434,8 @@ func (m initModel) formView() string {
found = strings.Join(m.found, ", ")
}
b.WriteString(" " + sDim.Render("Already in the certificate: "+found+". Add a DNS name or NAT address here.") + "\n\n")
fmt.Fprintf(&b, "%s\n %s\n", label(fWebhooks, "Notification webhooks"), m.webhooks.View())
b.WriteString(" " + sDim.Render("Where a parked model, a lost node or a spent budget is said: Slack, Mattermost, or any URL taking JSON. Comma-separated; a test is sent.") + "\n\n")
b.WriteString(check(fPostgres, "PostgreSQL", "rather than SQLite; installed here, for several control planes or a larger store"))
if m.in.mirrorURL != "" {
b.WriteString(" Model mirror " + sDim.Render("at "+m.in.mirrorURL+", from --mirror") + "\n")
+25 -3
View File
@@ -20,6 +20,8 @@ func TestTheInitFormBecomesInitsSettings(t *testing.T) {
m = typeIn(m, "ops@example.com")
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
m = typeIn(m, "llm.example.com, 203.0.113.10")
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
m = typeIn(m, "https://hooks.slack.com/services/T1/B2/x, json:https://alerts.example.com/in")
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(" ")}) // PostgreSQL on
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab},
tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(" ")}) // starter model off
@@ -33,11 +35,30 @@ func TestTheInitFormBecomesInitsSettings(t *testing.T) {
if !in.postgres || in.noMirror || in.noNode || !in.noModel || !in.noGuardModel {
t.Errorf("toggles: postgres %v mirror %v node %v model %v", in.postgres, !in.noMirror, !in.noNode, !in.noModel)
}
if hooks, err := parseWebhooks(in.webhooks); err != nil || len(hooks) != 2 || hooks[0].Format != "slack" || hooks[1].Format != "json" {
t.Errorf("webhooks %q: %+v %v", in.webhooks, hooks, err)
}
if m.(initModel).form {
t.Error("the form stayed up after Set up")
}
}
func TestAWebhookThatIsNoAddressKeepsTheFormUp(t *testing.T) {
in := &serverInit{}
var m tea.Model = newInitModel(in, nil, true)
for range fWebhooks {
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
}
m = typeIn(m, "hooks.slack.com/services/x")
for range fStart - fWebhooks {
m = send(m, tea.KeyMsg{Type: tea.KeyTab})
}
m = send(m, tea.KeyMsg{Type: tea.KeyEnter})
if !m.(initModel).form || m.(initModel).focus != fWebhooks || !strings.Contains(m.View(), "http(s) address") {
t.Errorf("an address without its scheme was taken:\n%s", m.View())
}
}
func TestAnOfflineBundleIsGivenInTheFormWithItsDigest(t *testing.T) {
in := &serverInit{}
var m tea.Model = newInitModel(in, nil, true)
@@ -117,9 +138,10 @@ func TestWithAMirrorElsewhereTheFormDoesNotOfferOne(t *testing.T) {
if v := m.View(); !strings.Contains(v, "at http://10.0.0.150:8082, from --mirror") {
t.Errorf("the mirror in use is not shown:\n%s", v)
}
// Tab from the addresses goes past the mirror line to PostgreSQL, then
// the node.
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab})
// Tab from the addresses goes to the webhooks, past the mirror line to
// PostgreSQL, then the node.
m = send(m, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab}, tea.KeyMsg{Type: tea.KeyTab},
tea.KeyMsg{Type: tea.KeyTab})
if f := m.(initModel).focus; f != fNode {
t.Errorf("focus %d, want the node toggle %d", f, fNode)
}
+44 -6
View File
@@ -15,6 +15,7 @@ import (
"github.com/charmbracelet/lipgloss"
joinpkg "gitea.mynodes.xyz/ilker/llmgw/internal/join"
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
"gitea.mynodes.xyz/ilker/llmgw/internal/preflight"
@@ -150,8 +151,8 @@ type setupModel struct {
ips []string
choice int // 0 one machine, 1 high availability
admin, names textField
afocus int // admin, names, next
admin, names, webhooks textField
afocus int // admin, names, webhooks, next
db [2]dbMachine
mfocus int // over both machines' fields, then the check button
@@ -184,8 +185,10 @@ func newSetupModel(host string, ips []string) setupModel {
admin.Placeholder, admin.CharLimit, admin.Width = "you@example.com", 200, 40
names := newTextField()
names.Placeholder, names.CharLimit, names.Width = "llm.example.com, 203.0.113.10", 400, 40
webhooks := newTextField()
webhooks.Placeholder, webhooks.CharLimit, webhooks.Width = "https://hooks.slack.com/services/…", 800, 40
return setupModel{
ref: &progRef{}, host: host, ips: ips, admin: admin, names: names,
ref: &progRef{}, host: host, ips: ips, admin: admin, names: names, webhooks: webhooks,
db: [2]dbMachine{newDBMachine(1), newDBMachine(2)}, spin: newSpinner(), width: 100, height: 30,
log: &[]string{}, logMu: &sync.Mutex{},
}
@@ -308,14 +311,17 @@ func (m setupModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
move := func(d int) {
m.afocus = (m.afocus + d + 3) % 3
m.afocus = (m.afocus + d + 4) % 4
m.admin.Blur()
m.names.Blur()
m.webhooks.Blur()
switch m.afocus {
case 0:
m.admin.Focus()
case 1:
m.names.Focus()
case 2:
m.webhooks.Focus()
}
}
switch k.Type {
@@ -326,7 +332,7 @@ func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
move(-1)
return m, nil
case tea.KeyEnter:
if m.afocus < 2 {
if m.afocus < 3 {
move(1)
return m, nil
}
@@ -336,6 +342,12 @@ func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
move(0)
return m, nil
}
if _, err := parseWebhooks(m.webhookList()); err != nil {
m.formErr = err.Error()
m.afocus = 2
move(0)
return m, nil
}
m.formErr = ""
m.page, m.mfocus = pMachines, 0
m.db[0].addr.Focus()
@@ -347,10 +359,23 @@ func (m setupModel) adminKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
m.admin, cmd = m.admin.Update(k)
case 1:
m.names, cmd = m.names.Update(k)
case 2:
m.webhooks, cmd = m.webhooks.Update(k)
}
return m, cmd
}
// webhookList is the webhooks asked for, one an item.
func (m setupModel) webhookList() []string {
var out []string
for _, s := range strings.Split(m.webhooks.Value(), ",") {
if s = strings.TrimSpace(s); s != "" {
out = append(out, s)
}
}
return out
}
func (m *setupModel) focusMachine() {
for i := range m.db {
for f := 0; f < mfCount; f++ {
@@ -589,6 +614,9 @@ func (m setupModel) command() string {
for _, t := range m.targets() {
parts = append(parts, "--patroni "+t)
}
for _, w := range m.webhookList() {
parts = append(parts, "--webhook "+w)
}
parts = append(parts, "--no-node", "--no-model")
return strings.Join(parts, " \\\n ")
}
@@ -600,6 +628,7 @@ func (m *setupModel) startRun() tea.Cmd {
in.admin = strings.TrimSpace(m.admin.Value())
in.hosts = m.namesList()
in.patroni = m.targets()
in.webhooks = m.webhookList()
// The control plane of a replicated setup is not where models run;
// nodes are added once it is up.
in.noNode, in.noModel = true, true
@@ -759,7 +788,9 @@ func (m setupModel) adminView() string {
sDim.Render("A password is generated and shown once, at the end."))
fmt.Fprintf(&b, "%s\n %s\n %s\n\n", label(1, "Other names nodes will reach this machine by"), m.names.View(),
sDim.Render("Known: "+strings.Join(append([]string{m.host}, m.ips...), ", ")+". Add a DNS name or a floating address."))
b.WriteString(button(m.afocus == 2, "Next") + "\n")
fmt.Fprintf(&b, "%s\n %s\n %s\n\n", label(2, "Notification webhooks"), m.webhooks.View(),
sDim.Render("Where a parked model, a lost node or a spent budget is said: Slack, Mattermost, or any URL taking JSON. Optional; a test is sent."))
b.WriteString(button(m.afocus == 3, "Next") + "\n")
return b.String()
}
@@ -880,6 +911,13 @@ func (m setupModel) planView() string {
fmt.Fprintf(&b, " %s %s\n %s\n", lipgloss.NewStyle().Bold(true).Render(m.host), sDim.Render(m.thisAddr()),
"control plane and gateway, model mirror, etcd (the third vote)")
fmt.Fprintf(&b, " administrator %s\n", strings.TrimSpace(m.admin.Value()))
if hooks, _ := parseWebhooks(m.webhookList()); len(hooks) > 0 {
var where []string
for _, h := range hooks {
where = append(where, notify.Redact(h.URL)+" ("+h.Format+")")
}
fmt.Fprintf(&b, " alerts to %s\n", strings.Join(where, ", "))
}
granted := false
for _, d := range m.db {
addr := strings.TrimSpace(d.addr.Value())
+2 -2
View File
@@ -45,7 +45,7 @@ func TestTheSetupAsksForWhatHighAvailabilityNeedsAndRefusesWhatWillNotDo(t *test
t.Fatalf("high availability did not lead to the administrator: page %d", s.page)
}
// No email: refused, with a reason.
m = press(m, "tab", "tab", "enter")
m = press(m, "tab", "tab", "tab", "enter")
if s := m.(setupModel); s.page != pAdmin || !strings.Contains(s.formErr, "email") {
t.Fatalf("went on without an administrator: %+v", s.formErr)
}
@@ -54,7 +54,7 @@ func TestTheSetupAsksForWhatHighAvailabilityNeedsAndRefusesWhatWillNotDo(t *test
if !strings.Contains(m.View(), "Not yet") {
t.Error("the refusal was not shown as a dialog")
}
m = press(m, "x", "a@example.com", "tab", "tab", "enter")
m = press(m, "x", "a@example.com", "tab", "tab", "tab", "enter")
s := m.(setupModel)
if s.page != pMachines {
t.Fatalf("page %d: %s", s.page, s.formErr)
+1 -1
View File
@@ -107,7 +107,7 @@ var Routes = []Route{
{Method: "GET", Path: "/v1/memory/{key}", Summary: "A note", Response: MemoryNote{}},
{Method: "GET", Path: "/v1/memory", Summary: "The project's notes", Response: MemoryList{}},
{Method: "DELETE", Path: "/v1/memory/{key}", Summary: "Forget a note"},
{Method: "POST", Path: "/mcp", Summary: "The gateway's MCP server: the memory notes, and pictures, 3D models, vision and workflows, as tools",
{Method: "POST", Path: "/mcp", Summary: "The gateway's MCP server: the memory notes as tools",
Request: MCPRequest{}, Response: anyJSON},
{Method: "GET", Path: "/me", Summary: "The caller: its key, its place in the tree, its limits and what it has used",
Response: Me{}},
+5 -1
View File
@@ -258,7 +258,11 @@ type Result struct {
// redact keeps a webhook's host and drops its path, where Slack and others
// put the secret.
func redact(u string) string {
func redact(u string) string { return Redact(u) }
// Redact is a webhook's address as it may be shown: its host, not the path
// that is its secret.
func Redact(u string) string {
if i := strings.Index(u, "://"); i >= 0 {
if j := strings.IndexByte(u[i+3:], '/'); j >= 0 {
return u[:i+3+j] + "/…"
+2 -6
View File
@@ -359,9 +359,7 @@ func (r *Router) Handler() http.Handler {
}
if r.auth == nil {
h := r.refuseWhileStopped(r.refuseLoops(mux))
r.selfHandler.Store(&h)
return h
return r.refuseWhileStopped(r.refuseLoops(mux))
}
exempt := []string{"/health", "/ready", "/leader", "/v1/openapi.json", "/", "/about", "/_static/", "/_lang/", "/claim/",
"/signup", "/login", "/login/", "/logout", "/profile", "/profile/", "/scim/",
@@ -375,9 +373,7 @@ func (r *Router) Handler() http.Handler {
// Wrapping the whole mux rather than checking inside each handler: adding a
// route must not be able to add an unauthenticated one by omission. Health
// stays open because a load balancer has no key.
h := r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
r.selfHandler.Store(&h)
return h
return r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
}
// refuseWhileStopped answers every request for a model with 503 while the
-937
View File
@@ -1,937 +0,0 @@
package router
import (
"bytes"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"image"
"image/color"
_ "image/jpeg" // a picture sent to look at may be a JPEG
"image/png"
"net/http"
"slices"
"strconv"
"strings"
"sync"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)
// --- MCP: pictures, 3D models, looking at pictures, workflows ---
//
// The gateway's own models offered as tools, so an agent — Claude Code
// driving Blender, say — makes a concept, a mesh from it and has a render
// looked at without a tool server of its own. Each tool makes its request
// of the gateway as the caller would, under the caller's key: the same
// allowlist, limits, licences left out, guardrails and accounting. What it
// makes is kept as a file of the caller's, said by id and by the address
// it is fetched from with the same key, since a tool server over HTTP
// cannot write to the caller's disk.
const (
// mcpBody bounds a tool call: up to four pictures to look at, as base64.
mcpBody = 64 << 20
// previewEdge is the long side of the picture shown to the agent beside
// the file: enough to see it, not a megapixel of tokens.
previewEdge = 512
// lookEdge is the long side of a picture sent to a vision model:
// 1280x720 is about 940 tokens.
lookEdge = 1280
// mcpWaitStreamed is how long a call answered as a stream waits for its
// job, saying how far along it is meanwhile; mcpWaitPlain how long one
// answered in one piece does, under the minute a proxy allows.
mcpWaitStreamed = 10 * time.Minute
mcpWaitPlain = 50 * time.Second
mcpWaitMax = 15 * time.Minute
mcpPoll = 2 * time.Second
)
// The prompt endings of generate_image's styles: what a picture a mesh is
// made from wants, and a surface that tiles.
var imageStyles = map[string]string{
"none": "",
"concept": ", single object centered, orthographic three-quarter view, flat even studio lighting, no shadows, " +
"no background elements, plain flat grey background, clean silhouette, game asset concept, high detail, no text, no watermark",
"side": ", single object, exact side view (profile), orthographic, flat even studio lighting, no shadows, " +
"plain flat grey background, clean silhouette, game asset concept, high detail, no text, no watermark",
"texture": ", seamless tileable texture, top-down orthographic, flat even lighting, no shadows, no perspective, no text",
}
// mcpMediaInstructions is said to an agent when it connects and the gateway
// has something to make pictures or 3D models with.
const mcpMediaInstructions = "Pictures, 3D models and a model that looks at pictures, on this gateway's own GPUs. " +
"A typical loop: generate_image (style concept) -> generate_mesh from its file_id -> import the GLB -> render -> " +
"review_images against the concept -> fix. Results are files: fetch them from their url with your API key. " +
"AI meshes suit rocks, debris and background props; hard-surface hero assets are usually better built by hand " +
"from the concept. A GPU holds one of these models at a time and switching costs a cold start, so batch calls " +
"of the same kind; gateway_status says which is loaded."
// mediaModels are the gateway's models by what they do, by name, deduplicated
// from the aliases that also resolve to them.
type mediaModels struct{ images, meshes, vision, workflows []string }
func (r *Router) mediaModels() mediaModels {
var out mediaModels
seen := map[string]bool{}
for _, m := range r.modelMap() {
if seen[m.Name] || m.IsGuard() {
continue
}
seen[m.Name] = true
if m.MakesImages() {
out.images = append(out.images, m.Name)
}
if m.MakesMeshes() {
out.meshes = append(out.meshes, m.Name)
}
if m.Capabilities.Vision {
out.vision = append(out.vision, m.Name)
}
if len(m.ComfyInputs()) > 0 {
out.workflows = append(out.workflows, m.Name)
}
}
for _, l := range [][]string{out.images, out.meshes, out.vision, out.workflows} {
slices.Sort(l)
}
return out
}
func (mm mediaModels) any() bool {
return len(mm.images)+len(mm.meshes)+len(mm.vision)+len(mm.workflows) > 0
}
// defaultOf is the model a tool uses when the call names none: the one the
// tool is best known with when the gateway has it, else the first.
func defaultOf(models []string, preferred ...string) string {
for _, p := range preferred {
if slices.Contains(models, p) {
return p
}
}
if len(models) == 0 {
return ""
}
return models[0]
}
func enumOrNone(models []string) map[string]any {
s := map[string]any{"type": "string"}
if len(models) > 0 {
s["enum"] = models
}
return s
}
// mediaTools are the tools for what the gateway's models can do: none for
// what it has no model for.
func (r *Router) mediaTools() []map[string]any {
mm := r.mediaModels()
var tools []map[string]any
fileArg := map[string]any{"type": "string",
"description": "A picture: a /v1/files id (file-…), a data URL, or base64"}
wait := map[string]any{"type": "integer", "minimum": 0, "maximum": int(mcpWaitMax.Seconds()),
"description": "Most seconds to wait for it before answering with the job to ask after (job_status)"}
if len(mm.images) > 0 {
tools = append(tools, map[string]any{
"name": "generate_image",
"description": "Make a picture from a prompt (~15 s warm, longer when the GPU must switch models). " +
"Keeps the PNG as a file and answers its file_id, url and seed, with a small preview. " +
"style: 'concept' (one object, 3/4 view, grey background: good for generate_mesh), 'side' " +
"(exact side view: best meshes for vehicles), 'texture' (seamless tileable surface), 'none' (the prompt as is).",
"inputSchema": map[string]any{
"type": "object",
"required": []string{"prompt"},
"properties": map[string]any{
"prompt": map[string]any{"type": "string"},
"style": map[string]any{"type": "string", "enum": []string{"none", "concept", "side", "texture"}, "default": "none"},
"size": map[string]any{"type": "integer", "description": "Square side in px, a multiple of 16, 256-2048", "default": 1024},
"seed": map[string]any{"type": "integer"},
"negative_prompt": map[string]any{"type": "string"},
"model": enumOrNone(mm.images),
},
},
})
}
if len(mm.meshes) > 0 {
tools = append(tools, map[string]any{
"name": "generate_mesh",
"description": "Make a 3D model (GLB) from a picture of one object on a plain background — the gateway " +
"cuts the background out itself — or from a prompt, for which a concept picture is made first. " +
"Takes 30 s to a few minutes; waits for it, saying how far along it is when the client takes a stream. " +
"Answers the GLB's file_id and url, and its licence. The mesh is about 1 m across, its orientation arbitrary.",
"inputSchema": map[string]any{
"type": "object",
"properties": map[string]any{
"image": fileArg,
"prompt": map[string]any{"type": "string", "description": "Used when no image is given"},
"model": enumOrNone(mm.meshes),
"seed": map[string]any{"type": "integer"},
"resolution": map[string]any{"type": "integer", "description": "The model's own grid, a multiple of 32 (TripoSR: 256 is enough)"},
"texture_resolution": map[string]any{"type": "integer", "minimum": 512, "maximum": 2048,
"description": "For a model that textures its mesh (sf3d): the texture's side in px"},
"inputs": map[string]any{"type": "object",
"description": "Any other input the model's workflow declares (gateway_status lists them): remesh, vertex_count, foreground_ratio, albedo_match…"},
"wait_seconds": wait,
},
},
})
}
if len(mm.vision) > 0 {
tools = append(tools, map[string]any{
"name": "review_images",
"description": "Ask a vision model about 1-4 pictures — a render beside its concept, say. A cheap critic " +
"for render-and-fix loops. json_schema makes it answer strict JSON. A cold start can take minutes " +
"when the GPU was making pictures.",
"inputSchema": map[string]any{
"type": "object",
"required": []string{"images", "question"},
"properties": map[string]any{
"images": map[string]any{"type": "array", "items": fileArg, "minItems": 1, "maxItems": 4},
"question": map[string]any{"type": "string"},
"json_schema": map[string]any{"type": "object", "description": "JSON Schema the answer must follow"},
"model": enumOrNone(mm.vision),
},
},
})
}
if len(mm.workflows) > 0 {
tools = append(tools, map[string]any{
"name": "run_workflow",
"description": "Run a ComfyUI workflow the gateway serves, with the inputs its catalog entry declares " +
"(gateway_status lists them). A file input is a /v1/files id, a data URL or base64. Answers its result file.",
"inputSchema": map[string]any{
"type": "object",
"required": []string{"model", "inputs"},
"properties": map[string]any{
"model": enumOrNone(mm.workflows),
"inputs": map[string]any{"type": "object"},
"wait_seconds": wait,
},
},
})
}
if mm.any() {
tools = append(tools,
map[string]any{
"name": "job_status",
"description": "Where a 3D model or workflow run stands, and its result file when done. wait_seconds waits for it to finish.",
"inputSchema": map[string]any{
"type": "object",
"required": []string{"job_id"},
"properties": map[string]any{"job_id": map[string]any{"type": "string"}, "wait_seconds": wait},
},
},
map[string]any{
"name": "gateway_status",
"description": "The gateway's picture, 3D, vision and workflow models: whether each is loaded on a GPU now, " +
"and a workflow's inputs; and your last few jobs. Use it to batch calls and avoid cold starts.",
"inputSchema": map[string]any{"type": "object", "properties": map[string]any{}},
})
}
return tools
}
// mediaCall is one tool call: who made it, and how to tell them how far it
// has got.
type mediaCall struct {
req *http.Request
progress func(percent float64, message string)
// streamed is a call whose answer is a stream, which may wait longer.
streamed bool
}
func (c mediaCall) wait(asked *int) time.Duration {
d := mcpWaitPlain
if c.streamed {
d = mcpWaitStreamed
}
if asked != nil {
d = min(time.Duration(*asked)*time.Second, mcpWaitMax)
}
return d
}
// isMediaTool is whether a tool is one of these.
func isMediaTool(name string) bool {
switch name {
case "generate_image", "generate_mesh", "review_images", "run_workflow", "job_status", "gateway_status":
return true
}
return false
}
// callMediaTool runs one of the tools. As with the notes, a failure is a
// tool result with isError set, which the model gets to read.
func (r *Router) callMediaTool(c mediaCall, name string, args json.RawMessage) map[string]any {
var res map[string]any
var err error
switch name {
case "generate_image":
res, err = r.toolGenerateImage(c, args)
case "generate_mesh":
res, err = r.toolGenerateMesh(c, args)
case "review_images":
res, err = r.toolReviewImages(c, args)
case "run_workflow":
res, err = r.toolRunWorkflow(c, args)
case "job_status":
res, err = r.toolJobStatus(c, args)
case "gateway_status":
res, err = r.toolGatewayStatus(c)
default:
err = errors.New("unknown tool " + name)
}
if err != nil {
return toolText(err.Error(), true)
}
return res
}
// self makes a request of the gateway's client API as the caller: their key,
// their agent name, and the context the call was made in.
func (r *Router) self(c mediaCall, method, path string, body any) (int, http.Header, []byte, error) {
h := r.selfHandler.Load()
if h == nil {
return 0, nil, nil, errors.New("the gateway's API is not mounted here")
}
var rd *bytes.Reader
if body != nil {
b, err := json.Marshal(body)
if err != nil {
return 0, nil, nil, err
}
rd = bytes.NewReader(b)
} else {
rd = bytes.NewReader(nil)
}
out, err := http.NewRequestWithContext(c.req.Context(), method, path, rd)
if err != nil {
return 0, nil, nil, err
}
for _, k := range []string{"Authorization", "X-Api-Key", agentHeader, "X-Forwarded-For", "X-Real-Ip"} {
if v := c.req.Header.Get(k); v != "" {
out.Header.Set(k, v)
}
}
out.RemoteAddr = c.req.RemoteAddr
if body != nil {
out.Header.Set("Content-Type", "application/json")
}
// A picture or a mesh comes back whole: past the 1 MiB a judge's
// answer is held to.
w := &bufferWriter{header: http.Header{}, max: mcpBody * 2}
(*h).ServeHTTP(w, out)
status := w.status
if status == 0 {
status = http.StatusOK
}
return status, w.header, w.buf.Bytes(), nil
}
// selfJSON is self for an answer in JSON, a failure said as the API said it.
func (r *Router) selfJSON(c mediaCall, method, path string, body, into any) (http.Header, error) {
status, h, b, err := r.self(c, method, path, body)
if err != nil {
return nil, err
}
if status >= 300 {
return h, fmt.Errorf("%s %s: %d %s", method, path, status, errorMessage(b))
}
if into == nil {
return h, nil
}
return h, json.Unmarshal(b, into)
}
// fileURL is where a file is fetched from, with the caller's key, as the
// caller reached the gateway.
func fileURL(req *http.Request, id string) string {
scheme := "http"
if req.TLS != nil {
scheme = "https"
}
if p := req.Header.Get("X-Forwarded-Proto"); p != "" {
scheme = strings.TrimSpace(strings.Split(p, ",")[0])
}
host := req.Host
if h := req.Header.Get("X-Forwarded-Host"); h != "" {
host = strings.TrimSpace(strings.Split(h, ",")[0])
}
return scheme + "://" + host + "/v1/files/" + id + "/content"
}
// keepFile keeps what a tool made as the caller's file, for as long as a
// job's result is kept.
func (r *Router) keepFile(c mediaCall, name string, content []byte) (store.File, error) {
p, ok := auth.FromContext(c.req.Context())
if !ok || r.store == nil {
return store.File{}, errors.New("this gateway keeps no files (it runs without authentication or a store)")
}
if r.noStore(callerSubject(c.req)) {
return store.File{}, errors.New("the gateway's policy for this caller keeps nothing, so no file is kept")
}
_, _, _, days := r.batchCfg()
now := time.Now()
f := store.File{ID: newID("file-"), Owner: ownerOf(p), Purpose: purposeJobOutput, Filename: name,
Bytes: int64(len(content)), CreatedAt: now, ExpiresAt: now.AddDate(0, 0, days)}
return f, r.store.PutFile(c.req.Context(), f, content)
}
// picture reads a picture argument: a file of the caller's by id, a data URL,
// or base64.
func (r *Router) picture(c mediaCall, s string) ([]byte, error) {
s = strings.TrimSpace(s)
switch {
case s == "":
return nil, errors.New("no picture given")
case strings.HasPrefix(s, "file-"):
p, ok := auth.FromContext(c.req.Context())
if !ok || r.store == nil {
return nil, errors.New("this gateway keeps no files: send the picture as base64 or a data URL")
}
f, err := r.store.FileByID(c.req.Context(), s)
if errors.Is(err, store.ErrNotFound) || err == nil && f.Owner != ownerOf(p) {
return nil, fmt.Errorf("no file %s of yours", s)
}
if err != nil {
return nil, err
}
return r.store.FileContent(c.req.Context(), s)
case strings.HasPrefix(s, "data:"):
_, after, ok := strings.Cut(s, ";base64,")
if !ok {
return nil, errors.New("a data URL without base64 in it")
}
s = after
}
b, err := base64.StdEncoding.DecodeString(s)
if err != nil {
return nil, errors.New("the picture is neither a file id, a data URL nor base64")
}
return b, nil
}
func dataURL(png []byte) string {
return "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)
}
// shrink returns a picture as a PNG whose long side is at most edge.
func shrink(pic []byte, edge int) ([]byte, error) {
src, format, err := image.Decode(bytes.NewReader(pic))
if err != nil {
return nil, fmt.Errorf("not a picture this gateway reads (PNG or JPEG): %w", err)
}
b := src.Bounds()
w, h := b.Dx(), b.Dy()
if max(w, h) <= edge {
if format == "png" {
return pic, nil
}
edge = max(w, h) // only re-encoded
}
scale := float64(edge) / float64(max(w, h))
dw, dh := max(1, int(float64(w)*scale+0.5)), max(1, int(float64(h)*scale+0.5))
dst := image.NewNRGBA(image.Rect(0, 0, dw, dh))
// Each pixel the average of the box of pixels it stands for.
for y := range dh {
y0, y1 := b.Min.Y+y*h/dh, b.Min.Y+max((y+1)*h/dh, y*h/dh+1)
for x := range dw {
x0, x1 := b.Min.X+x*w/dw, b.Min.X+max((x+1)*w/dw, x*w/dw+1)
var sr, sg, sb, sa, n uint64
for yy := y0; yy < y1; yy++ {
for xx := x0; xx < x1; xx++ {
cr, cg, cb, ca := src.At(xx, yy).RGBA()
sr, sg, sb, sa, n = sr+uint64(cr), sg+uint64(cg), sb+uint64(cb), sa+uint64(ca), n+1
}
}
// Premultiplied sums, back to straight colour.
if sa == 0 {
continue
}
dst.SetNRGBA(x, y, color.NRGBA{
R: uint8(sr * 0xff / sa), G: uint8(sg * 0xff / sa), B: uint8(sb * 0xff / sa), A: uint8(sa / n >> 8),
})
}
}
var out bytes.Buffer
if err := png.Encode(&out, dst); err != nil {
return nil, err
}
return out.Bytes(), nil
}
// --- the tools ---
func (r *Router) toolGenerateImage(c mediaCall, raw json.RawMessage) (map[string]any, error) {
var a struct {
Prompt string `json:"prompt"`
Style string `json:"style"`
NegativePrompt string `json:"negative_prompt"`
Model string `json:"model"`
Size int `json:"size"`
Seed *int64 `json:"seed"`
}
if err := decodeArgs(raw, &a); err != nil {
return nil, err
}
res, _, err := r.makeImage(c, a.Prompt, a.Style, a.NegativePrompt, a.Model, a.Size, a.Seed)
if err != nil {
return nil, err
}
return res, nil
}
// makeImage makes a picture and keeps it: the tool's answer, and the PNG.
func (r *Router) makeImage(c mediaCall, prompt, style, negative, model string, size int, seed *int64) (map[string]any, []byte, error) {
if strings.TrimSpace(prompt) == "" {
return nil, nil, errors.New("prompt is required")
}
if style == "" {
style = "none"
}
suffix, ok := imageStyles[style]
if !ok {
return nil, nil, fmt.Errorf("style %q: concept, side, texture or none", style)
}
if model == "" {
model = defaultOf(r.mediaModels().images, "flux-schnell")
}
if size == 0 {
size = 1024
}
body := map[string]any{"model": model, "prompt": prompt + suffix, "n": 1,
"size": fmt.Sprintf("%dx%d", size, size), "response_format": "b64_json"}
if seed != nil {
body["seed"] = *seed
}
if negative != "" {
body["negative_prompt"] = negative
}
c.progress(5, "making the picture with "+model)
var out struct {
Data []struct {
B64 string `json:"b64_json"`
} `json:"data"`
}
h, err := r.selfJSON(c, http.MethodPost, imagesPath, body, &out)
if err != nil {
return nil, nil, err
}
if len(out.Data) == 0 {
return nil, nil, errors.New(model + " made no picture")
}
pic, err := base64.StdEncoding.DecodeString(out.Data[0].B64)
if err != nil {
return nil, nil, err
}
info := map[string]any{"model": model, "width": size, "height": size}
if s, err := strconv.ParseInt(h.Get("x-llmgw-seed"), 10, 64); err == nil {
info["seed"] = s
}
if m, ok := r.modelMap()[model]; ok && m.License != nil {
info["license"] = m.License
}
content := []map[string]any{}
if f, err := r.keepFile(c, "image-"+strings.TrimPrefix(newID(""), "-")+".png", pic); err == nil {
info["file_id"], info["url"], info["bytes"] = f.ID, fileURL(c.req, f.ID), f.Bytes
} else {
// Nowhere to keep it: the picture itself, whole.
info["note"] = err.Error() + "; the picture is in this answer at full size"
content = append(content, map[string]any{"type": "image", "mimeType": "image/png", "data": out.Data[0].B64})
}
text, _ := json.Marshal(info)
content = append([]map[string]any{{"type": "text", "text": string(text)}}, content...)
if _, kept := info["file_id"]; kept {
if small, err := shrink(pic, previewEdge); err == nil {
content = append(content, map[string]any{"type": "image", "mimeType": "image/png",
"data": base64.StdEncoding.EncodeToString(small)})
}
}
return map[string]any{"content": content, "isError": false}, pic, nil
}
func (r *Router) toolGenerateMesh(c mediaCall, raw json.RawMessage) (map[string]any, error) {
var a struct {
Image, Prompt, Model string
Seed *int64
Resolution int
TextureResolution int `json:"texture_resolution"`
Inputs map[string]any
WaitSeconds *int `json:"wait_seconds"`
}
if err := decodeArgs(raw, &a); err != nil {
return nil, err
}
mm := r.mediaModels()
if a.Model == "" {
a.Model = defaultOf(mm.meshes, "sf3d", "triposr")
}
m, ok := r.modelMap()[a.Model]
if !ok || !m.MakesMeshes() {
return nil, fmt.Errorf("%q makes no 3D models here; these do: %s", a.Model, strings.Join(mm.meshes, ", "))
}
extra := map[string]any{}
var pic []byte
var err error
switch {
case a.Image != "":
if pic, err = r.picture(c, a.Image); err != nil {
return nil, err
}
case a.Prompt != "":
res, made, err := r.makeImage(c, a.Prompt, "concept", "", "", 0, nil)
if err != nil {
return nil, fmt.Errorf("the concept picture: %w", err)
}
pic = made
var info map[string]any
if json.Unmarshal([]byte(res["content"].([]map[string]any)[0]["text"].(string)), &info) == nil {
extra["concept_file_id"], extra["concept_seed"] = info["file_id"], info["seed"]
}
default:
return nil, errors.New("an image, or a prompt to make one from")
}
// A model whose workflow declares its inputs takes them at the
// workflows' path, every one it names; one that does not, the 3D path's
// seed and grid.
path, body := meshesPath, map[string]any{"model": a.Model, "image": dataURL(pic)}
if declared := m.ComfyInputs(); len(declared) > 0 {
names := map[string]string{}
var image string
for _, in := range declared {
names[in.Name] = in.Type
if image == "" && in.Type == "image" {
image = in.Name
}
}
if image == "" {
return nil, fmt.Errorf("%s's workflow takes no picture", a.Model)
}
inputs := map[string]any{image: dataURL(pic)}
for k, v := range a.Inputs {
inputs[k] = v
}
// A textured mesh's resolution is its texture's; a grid's, the grid's.
if _, ok := names["resolution"]; ok {
switch {
case a.TextureResolution > 0:
inputs["resolution"] = a.TextureResolution
case a.Resolution > 0:
inputs["resolution"] = a.Resolution
}
}
path, body = workflowsPath, map[string]any{"model": a.Model, "inputs": inputs}
if a.Seed != nil {
body["seed"] = *a.Seed
}
} else {
if a.Seed != nil {
body["seed"] = *a.Seed
}
if a.Resolution > 0 {
body["resolution"] = a.Resolution
}
}
var job struct {
ID string `json:"id"`
}
if _, err := r.selfJSON(c, http.MethodPost, path, body, &job); err != nil {
return nil, err
}
c.progress(10, "queued as "+job.ID)
return r.waitJob(c, job.ID, c.wait(a.WaitSeconds), extra)
}
func (r *Router) toolRunWorkflow(c mediaCall, raw json.RawMessage) (map[string]any, error) {
var a struct {
Model string
Inputs map[string]any
WaitSeconds *int `json:"wait_seconds"`
}
if err := decodeArgs(raw, &a); err != nil {
return nil, err
}
m, ok := r.modelMap()[a.Model]
if !ok || len(m.ComfyInputs()) == 0 {
return nil, fmt.Errorf("%q runs no workflow here; these do: %s", a.Model, strings.Join(r.mediaModels().workflows, ", "))
}
// A file input given as a file id goes as it is: the API takes those.
var job struct {
ID string `json:"id"`
}
if _, err := r.selfJSON(c, http.MethodPost, workflowsPath, map[string]any{"model": a.Model, "inputs": a.Inputs}, &job); err != nil {
return nil, err
}
c.progress(10, "queued as "+job.ID)
return r.waitJob(c, job.ID, c.wait(a.WaitSeconds), nil)
}
func (r *Router) toolJobStatus(c mediaCall, raw json.RawMessage) (map[string]any, error) {
var a struct {
JobID string `json:"job_id"`
WaitSeconds *int `json:"wait_seconds"`
}
if err := decodeArgs(raw, &a); err != nil {
return nil, err
}
if a.JobID == "" {
return nil, errors.New("job_id is required")
}
d := time.Duration(0)
if a.WaitSeconds != nil {
d = c.wait(a.WaitSeconds)
}
return r.waitJob(c, a.JobID, d, nil)
}
// jobView is the part of a job a tool says.
type jobView struct {
ID string `json:"id"`
Model string `json:"model"`
Status string `json:"status"`
Progress json.RawMessage `json:"progress"`
OutputFileID *string `json:"output_file_id"`
ContentType *string `json:"content_type"`
Error *struct {
Message string `json:"message"`
} `json:"error"`
License json.RawMessage `json:"license"`
}
// waitJob asks after a job until it is done or d has passed, saying how far
// along it is, and answers where it stands: its file when done.
func (r *Router) waitJob(c mediaCall, id string, d time.Duration, extra map[string]any) (map[string]any, error) {
deadline := time.Now().Add(d)
var j jobView
for {
if _, err := r.selfJSON(c, http.MethodGet, workflowsPath+"/"+id, nil, &j); err != nil {
return nil, err
}
if store.JobDone(j.Status) || !time.Now().Before(deadline) {
break
}
pct, msg := jobProgress(j)
c.progress(pct, msg)
select {
case <-c.req.Context().Done():
return nil, c.req.Context().Err()
case <-time.After(mcpPoll):
}
}
info := map[string]any{"job_id": j.ID, "model": j.Model, "status": j.Status}
for k, v := range extra {
if v != nil {
info[k] = v
}
}
if len(j.License) > 0 && string(j.License) != "null" {
info["license"] = j.License
}
switch {
case j.Status == store.JobCompleted && j.OutputFileID != nil:
info["file_id"], info["url"] = *j.OutputFileID, fileURL(c.req, *j.OutputFileID)
if j.ContentType != nil {
info["content_type"] = *j.ContentType
}
if r.store != nil {
if f, err := r.store.FileByID(c.req.Context(), *j.OutputFileID); err == nil {
info["bytes"] = f.Bytes
}
}
case store.JobDone(j.Status):
if j.Error != nil {
info["error"] = j.Error.Message
}
b, _ := json.Marshal(info)
return toolText(string(b), true), nil
default:
info["progress"] = j.Progress
info["next"] = "not done yet: call job_status with this job_id (and wait_seconds) to wait for it"
}
b, _ := json.Marshal(info)
return toolText(string(b), false), nil
}
// jobProgress reads a job's progress as a percent and a phrase.
func jobProgress(j jobView) (float64, string) {
var p struct {
State string `json:"state"`
Ahead int `json:"ahead"`
Percent float64 `json:"percent"`
Node string `json:"node_class"`
HeldBy string `json:"held_by"`
}
_ = json.Unmarshal(j.Progress, &p)
switch {
case p.State == "queued":
return 10, fmt.Sprintf("%s: queued, %d ahead", j.ID, p.Ahead)
case p.State == "waiting_gpu":
return 10, j.ID + ": waiting for the GPU, held by " + p.HeldBy
case p.State == "running":
return 10 + 0.9*p.Percent, fmt.Sprintf("%s: %s %.0f%%", j.ID, p.Node, p.Percent)
}
return 10, j.ID + ": " + j.Status
}
func (r *Router) toolReviewImages(c mediaCall, raw json.RawMessage) (map[string]any, error) {
var a struct {
Images []string
Question string
JSONSchema map[string]any `json:"json_schema"`
Model string
}
if err := decodeArgs(raw, &a); err != nil {
return nil, err
}
if len(a.Images) == 0 || len(a.Images) > 4 {
return nil, errors.New("one to four images")
}
if strings.TrimSpace(a.Question) == "" {
return nil, errors.New("question is required")
}
if a.Model == "" {
a.Model = defaultOf(r.mediaModels().vision, "vlm")
}
if a.Model == "" {
return nil, errors.New("no model here looks at pictures")
}
parts := []map[string]any{{"type": "text", "text": a.Question}}
for i, s := range a.Images {
pic, err := r.picture(c, s)
if err != nil {
return nil, fmt.Errorf("image %d: %w", i+1, err)
}
small, err := shrink(pic, lookEdge)
if err != nil {
return nil, fmt.Errorf("image %d: %w", i+1, err)
}
parts = append(parts, map[string]any{"type": "image_url", "image_url": map[string]any{"url": dataURL(small)}})
}
body := map[string]any{"model": a.Model, "messages": []map[string]any{{"role": "user", "content": parts}}}
if a.JSONSchema != nil {
body["response_format"] = map[string]any{"type": "json_schema",
"json_schema": map[string]any{"name": "answer", "schema": a.JSONSchema, "strict": true}}
}
c.progress(10, "asking "+a.Model)
var out struct {
Choices []struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
} `json:"choices"`
}
if _, err := r.selfJSON(c, http.MethodPost, "/v1/chat/completions", body, &out); err != nil {
return nil, err
}
if len(out.Choices) == 0 {
return nil, errors.New(a.Model + " gave no answer")
}
return toolText(out.Choices[0].Message.Content, false), nil
}
func (r *Router) toolGatewayStatus(c mediaCall) (map[string]any, error) {
var models struct {
Data []map[string]any `json:"data"`
}
if _, err := r.selfJSON(c, http.MethodGet, "/v1/models", nil, &models); err != nil {
return nil, err
}
var listed []map[string]any
for _, m := range models.Data {
caps, _ := m["capabilities"].(map[string]any)
if caps["images"] != true && caps["meshes"] != true && caps["vision"] != true && m["workflow"] == nil {
continue
}
e := map[string]any{"id": m["id"], "loaded": m["loaded"], "capabilities": caps}
if w := m["workflow"]; w != nil {
e["workflow"] = w
}
if l := m["license"]; l != nil {
e["license"] = l
}
listed = append(listed, e)
}
out := map[string]any{"models": listed}
var jobs struct {
Data []jobView `json:"data"`
}
if _, err := r.selfJSON(c, http.MethodGet, workflowsPath+"?limit=5", nil, &jobs); err == nil {
recent := make([]map[string]any, 0, len(jobs.Data))
for _, j := range jobs.Data {
recent = append(recent, map[string]any{"id": j.ID, "model": j.Model, "status": j.Status})
}
out["recent_jobs"] = recent
}
return toolJSON(out), nil
}
// decodeArgs reads a tool's arguments, by their JSON names: a field's own
// name in lower case unless its tag says.
func decodeArgs(raw json.RawMessage, into any) error {
if len(raw) == 0 || string(raw) == "null" {
raw = []byte("{}")
}
if err := json.Unmarshal(raw, into); err != nil {
return fmt.Errorf("the arguments: %w", err)
}
return nil
}
// streamWriter answers a tool call as server-sent events: progress
// notifications while it runs, then the answer.
type streamWriter struct {
mu sync.Mutex
w http.ResponseWriter
token json.RawMessage
}
func (s *streamWriter) write(text string) {
s.mu.Lock()
defer s.mu.Unlock()
_, _ = fmt.Fprint(s.w, text)
if f, ok := s.w.(http.Flusher); ok {
f.Flush()
}
}
func (s *streamWriter) event(v any) {
b, _ := json.Marshal(v)
s.write("event: message\ndata: " + string(b) + "\n\n")
}
// keepAlive writes a comment every while until stop is closed: a picture
// whose model starts cold is minutes of nothing said, which a proxy reads as
// a connection gone idle.
func (s *streamWriter) keepAlive(stop <-chan struct{}, every time.Duration) {
t := time.NewTicker(every)
defer t.Stop()
for {
select {
case <-stop:
return
case <-t.C:
s.write(": still working\n\n")
}
}
}
func (s *streamWriter) progress(percent float64, message string) {
if len(s.token) == 0 {
return
}
s.event(map[string]any{"jsonrpc": "2.0", "method": "notifications/progress", "params": map[string]any{
"progressToken": s.token, "progress": percent, "total": 100, "message": message}})
}
// acceptsStream is a client that takes a POST's answer as a stream, as the
// streamable HTTP transport lets it.
func acceptsStream(req *http.Request) bool {
return strings.Contains(req.Header.Get("Accept"), "text/event-stream")
}
-266
View File
@@ -1,266 +0,0 @@
package router
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"image"
"image/png"
"net/http"
"net/http/httptest"
"strings"
"testing"
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)
// mcp makes one JSON-RPC call at /mcp and answers its result; streamed, as a
// client that takes server-sent events, it answers the events' data too.
func mcp(t *testing.T, h http.Handler, secret, method string, params any, streamed bool) (map[string]any, []string) {
t.Helper()
b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": method, "params": params})
req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(b))
req.Header.Set("Content-Type", "application/json")
if secret != "" {
req.Header.Set("Authorization", "Bearer "+secret)
}
if streamed {
req.Header.Set("Accept", "application/json, text/event-stream")
}
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("%s: %d %s", method, rec.Code, rec.Body.String())
}
body := rec.Body.String()
var events []string
if streamed {
if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" {
t.Fatalf("a streamed call answered %s: %s", ct, body)
}
for _, line := range strings.Split(body, "\n") {
if d, ok := strings.CutPrefix(line, "data: "); ok {
events = append(events, d)
}
}
if len(events) == 0 {
t.Fatalf("no events: %s", body)
}
body = events[len(events)-1]
}
var resp struct {
Result map[string]any `json:"result"`
Error any `json:"error"`
}
if err := json.Unmarshal([]byte(body), &resp); err != nil || resp.Error != nil {
t.Fatalf("%s: %v %s", method, err, body)
}
return resp.Result, events
}
func toolNames(res map[string]any) []string {
var names []string
for _, tl := range res["tools"].([]any) {
names = append(names, tl.(map[string]any)["name"].(string))
}
return names
}
// resultInfo is a tool result's first text, as the JSON it holds.
func resultInfo(t *testing.T, res map[string]any) map[string]any {
t.Helper()
if res["isError"] == true {
t.Fatalf("the tool failed: %v", res["content"])
}
var info map[string]any
text := res["content"].([]any)[0].(map[string]any)["text"].(string)
if err := json.Unmarshal([]byte(text), &info); err != nil {
t.Fatalf("not JSON: %s", text)
}
return info
}
func TestTheMCPToolsAreWhatTheGatewaysModelsCanDo(t *testing.T) {
c := newMeshComfy(t)
rt, secret := meshRouter(t, c)
h := rt.Handler()
init, _ := mcp(t, h, secret, "initialize", map[string]any{"protocolVersion": "2025-06-18"}, false)
if init["protocolVersion"] != "2025-06-18" || !strings.Contains(init["instructions"].(string), "batch calls") {
t.Errorf("initialize: %v", init)
}
list, _ := mcp(t, h, secret, "tools/list", map[string]any{}, false)
names := strings.Join(toolNames(list), ",")
for _, want := range []string{"remember", "generate_image", "generate_mesh", "job_status", "gateway_status"} {
if !strings.Contains(names, want) {
t.Errorf("no %s in %s", want, names)
}
}
// Nothing here looks at pictures or declares a workflow's inputs.
if strings.Contains(names, "review_images") || strings.Contains(names, "run_workflow") {
t.Errorf("tools for models the gateway has not: %s", names)
}
// A gateway with only text models offers the notes alone.
plain := newTestRouter(t, baseConfig()).Handler()
list, _ = mcp(t, plain, "", "tools/list", map[string]any{}, false)
if got := strings.Join(toolNames(list), ","); got != "remember,recall,forget" {
t.Errorf("a text-only gateway offers %s", got)
}
}
func TestAMeshIsMadeThroughMCPAndKeptAsTheCallersFile(t *testing.T) {
c := newMeshComfy(t)
rt, secret := meshRouter(t, c)
h := rt.Handler()
args := map[string]any{"model": "trellis2", "image": "data:image/png;base64," + testPicture(t), "seed": 3}
for _, streamed := range []bool{false, true} {
res, events := mcp(t, h, secret, "tools/call", map[string]any{"name": "generate_mesh", "arguments": args,
"_meta": map[string]any{"progressToken": "tok"}}, streamed)
info := resultInfo(t, res)
if info["status"] != "completed" || info["license"].(map[string]any)["name"] != "MIT" {
t.Fatalf("streamed %v: %v", streamed, info)
}
id, _ := info["file_id"].(string)
if !strings.HasSuffix(info["url"].(string), "/v1/files/"+id+"/content") {
t.Errorf("url %v", info["url"])
}
// Fetched with the same key, it is the mesh.
code, _, file := call(t, h, secret, http.MethodGet, "/v1/files/"+id+"/content", nil)
if code != http.StatusOK || !strings.HasPrefix(file, "glTF") {
t.Fatalf("the file is not the mesh: %d", code)
}
if streamed && !strings.Contains(strings.Join(events, "\n"), `"progressToken":"tok"`) {
t.Errorf("no progress said while it ran: %v", events)
}
}
// The job went through the client API as the caller: it is theirs.
_, list, _ := call(t, h, secret, http.MethodGet, meshesPath, nil)
if n := len(list["data"].([]any)); n != 2 {
t.Errorf("%d jobs of the caller's, want the two made", n)
}
// Under no key, nothing is made.
b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "tools/call",
"params": map[string]any{"name": "generate_mesh", "arguments": args}})
req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(b))
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusUnauthorized {
t.Errorf("a call with no key: %d", rec.Code)
}
}
func TestAPictureMadeThroughMCPComesWholeWhereNoFileCanBeKept(t *testing.T) {
rt, queued, mu := imageRouter(t)
h := rt.Handler()
res, _ := mcp(t, h, "", "tools/call", map[string]any{"name": "generate_image",
"arguments": map[string]any{"prompt": "a rock", "style": "concept", "size": 512, "seed": 9}}, false)
info := resultInfo(t, res)
if info["seed"] != 9.0 || info["model"] != "flux" || !strings.Contains(info["note"].(string), "keeps no files") {
t.Errorf("info %v", info)
}
content := res["content"].([]any)
if len(content) != 2 || content[1].(map[string]any)["type"] != "image" {
t.Fatalf("the picture is not in the answer: %v", content)
}
if b, _ := base64.StdEncoding.DecodeString(content[1].(map[string]any)["data"].(string)); string(b) != "png-bytes" {
t.Errorf("picture %q", b)
}
mu.Lock()
wf, _ := json.Marshal(*queued)
mu.Unlock()
if !strings.Contains(string(wf), "a rock, single object centered") || !strings.Contains(string(wf), `"width":512`) {
t.Errorf("the style was not added to the prompt: %s", wf)
}
}
func TestAPictureIsShrunkToItsLongSide(t *testing.T) {
var b bytes.Buffer
_ = png.Encode(&b, image.NewRGBA(image.Rect(0, 0, 2000, 1000)))
out, err := shrink(b.Bytes(), 512)
if err != nil {
t.Fatal(err)
}
im, err := png.Decode(bytes.NewReader(out))
if err != nil || im.Bounds().Dx() != 512 || im.Bounds().Dy() != 256 {
t.Fatalf("%v %v", im.Bounds(), err)
}
b.Reset()
_ = png.Encode(&b, image.NewRGBA(image.Rect(0, 0, 100, 50)))
if out, _ := shrink(b.Bytes(), 512); !bytes.Equal(out, b.Bytes()) {
t.Error("a small PNG was re-encoded")
}
}
func TestAWorkflowRunsThroughMCPFromItsDeclaredInputs(t *testing.T) {
wav := testWAV()
c := newRunComfy(t, wav, "song_00001_.wav")
rt, secret, _ := runRouter(t, c, songWorkflow, catalog.OutputAudio)
h := rt.Handler()
list, _ := mcp(t, h, secret, "tools/list", map[string]any{}, false)
if !strings.Contains(strings.Join(toolNames(list), ","), "run_workflow") {
t.Fatalf("no run_workflow for a model that declares its inputs: %v", toolNames(list))
}
res, _ := mcp(t, h, secret, "tools/call", map[string]any{"name": "run_workflow", "arguments": map[string]any{
"model": "ace", "inputs": map[string]any{"tags": "synthwave", "cover": "data:image/png;base64," + testPicture(t)}}}, false)
info := resultInfo(t, res)
if info["status"] != "completed" || info["content_type"] != "audio/wav" {
t.Fatalf("%v", info)
}
if _, _, file := call(t, h, secret, http.MethodGet, "/v1/files/"+info["file_id"].(string)+"/content", nil); file != string(wav) {
t.Error("the file is not the song")
}
// A run refused says why, as the API did.
res, _ = mcp(t, h, secret, "tools/call", map[string]any{"name": "run_workflow", "arguments": map[string]any{
"model": "ace", "inputs": map[string]any{"cover": "data:image/png;base64," + testPicture(t)}}}, false)
if res["isError"] != true || !strings.Contains(fmt.Sprint(res["content"]), "tags") {
t.Errorf("a run missing a required input: %v", res)
}
}
func TestPicturesAreLookedAtByAVisionModelShrunkAndToASchema(t *testing.T) {
e := newFakeEngine(t, "vlm")
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"vlm"}})
vlm := catModel("vlm", nil)
vlm.Capabilities.Vision = true
cfg.Models = []catalog.Model{vlm}
h := newTestRouter(t, cfg).Handler()
var big bytes.Buffer
_ = png.Encode(&big, image.NewRGBA(image.Rect(0, 0, 3000, 1500)))
res, _ := mcp(t, h, "", "tools/call", map[string]any{"name": "review_images", "arguments": map[string]any{
"images": []string{base64.StdEncoding.EncodeToString(big.Bytes())},
"question": "Does the render match the concept?",
"json_schema": map[string]any{"type": "object", "properties": map[string]any{"match": map[string]any{"type": "boolean"}}},
}}, false)
if res["isError"] == true || !strings.Contains(fmt.Sprint(res["content"]), "hi") {
t.Fatalf("%v", res)
}
var sent struct {
ResponseFormat struct {
Type string `json:"type"`
} `json:"response_format"`
Messages []struct {
Content []struct {
Type string `json:"type"`
ImageURL struct {
URL string `json:"url"`
} `json:"image_url"`
} `json:"content"`
} `json:"messages"`
}
if err := json.Unmarshal(e.last.Load().([]byte), &sent); err != nil {
t.Fatal(err)
}
if sent.ResponseFormat.Type != "json_schema" || len(sent.Messages) != 1 || len(sent.Messages[0].Content) != 2 {
t.Fatalf("sent %+v", sent)
}
b64, _ := strings.CutPrefix(sent.Messages[0].Content[1].ImageURL.URL, "data:image/png;base64,")
pic, _ := base64.StdEncoding.DecodeString(b64)
im, err := png.Decode(bytes.NewReader(pic))
if err != nil || im.Bounds().Dx() != lookEdge || im.Bounds().Dy() != lookEdge/2 {
t.Errorf("the picture sent is %v (%v), not shrunk to %d", im.Bounds(), err, lookEdge)
}
}
+6 -51
View File
@@ -5,7 +5,6 @@ import (
"errors"
"io"
"net/http"
"slices"
"strconv"
"time"
@@ -121,8 +120,7 @@ func (r *Router) handlePutNote(w http.ResponseWriter, req *http.Request) {
Text string `json:"text"`
Tags []string `json:"tags"`
}
// Pictures to look at come in the call, as base64.
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, mcpBody))
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
if err != nil || json.Unmarshal(raw, &body) != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error",
`body must be JSON like {"text": "...", "tags": ["..."]}`)
@@ -275,66 +273,23 @@ func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
resp := rpcResponse{JSONRPC: "2.0", ID: rpc.ID}
switch rpc.Method {
case "initialize":
result := map[string]any{
"protocolVersion": mcpVersionFor(rpc.Params),
resp.Result = map[string]any{
"protocolVersion": mcpProtocolVer,
"capabilities": map[string]any{"tools": map[string]any{}},
"serverInfo": map[string]any{"name": "llmgw", "version": "1"},
"serverInfo": map[string]any{"name": "llmgw-memory", "version": "1"},
}
if r.mediaModels().any() {
result["instructions"] = mcpMediaInstructions
}
resp.Result = result
case "ping":
resp.Result = map[string]any{}
case "tools/list":
resp.Result = map[string]any{"tools": append(slices.Clone(mcpTools), r.mediaTools()...)}
resp.Result = map[string]any{"tools": mcpTools}
case "tools/call":
var call struct {
Name string `json:"name"`
Arguments json.RawMessage `json:"arguments"`
Meta struct {
ProgressToken json.RawMessage `json:"progressToken"`
} `json:"_meta"`
}
_ = json.Unmarshal(rpc.Params, &call)
if !isMediaTool(call.Name) {
resp.Result = r.callMemoryTool(req, rpc.Params)
break
}
if acceptsStream(req) {
// Answered as a stream: how far along it is while it runs, which
// a mesh's minutes want, then the answer.
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.WriteHeader(http.StatusOK)
sw := &streamWriter{w: w, token: call.Meta.ProgressToken}
stop := make(chan struct{})
go sw.keepAlive(stop, 15*time.Second)
resp.Result = r.callMediaTool(mediaCall{req: req, progress: sw.progress, streamed: true}, call.Name, call.Arguments)
close(stop)
sw.event(resp)
return
}
resp.Result = r.callMediaTool(mediaCall{req: req, progress: func(float64, string) {}}, call.Name, call.Arguments)
resp.Result = r.callMemoryTool(req, rpc.Params)
default:
resp.Error = &rpcError{Code: -32601, Message: "method not found: " + rpc.Method}
}
writeJSON(w, http.StatusOK, resp)
}
// mcpVersionFor is the protocol version answered: the client's, of those
// this server speaks, else its own.
func mcpVersionFor(params json.RawMessage) string {
var p struct {
ProtocolVersion string `json:"protocolVersion"`
}
_ = json.Unmarshal(params, &p)
if slices.Contains([]string{"2025-03-26", "2025-06-18"}, p.ProtocolVersion) {
return p.ProtocolVersion
}
return mcpProtocolVer
}
func toolText(text string, isError bool) map[string]any {
return map[string]any{
"content": []map[string]any{{"type": "text", "text": text}},
-4
View File
@@ -51,10 +51,6 @@ type Router struct {
cfg config.Router
log *slog.Logger
// selfHandler is the client API as Handler built it, key check and all:
// what the MCP tools make their requests through, as the caller.
selfHandler atomic.Pointer[http.Handler]
reg *endpoint.Registry
poller *scrape.Poller
px *proxy.Proxy