Compare commits

...
Author SHA1 Message Date
Ilker ManapandClaude Opus 5.5 55d5e13534 Pictures, 3D models, a vision model and workflows as tools at /mcp
An agent driving Blender made its concepts, meshes and render reviews
through a tool server of the game's own, which called this gateway; the
gateway's MCP server offered only the memory notes. It now offers
generate_image, generate_mesh, review_images, run_workflow, job_status and
gateway_status beside them, each only where a model for it is here.

Each tool makes its request of the client API as the caller, through the
handler Handler built, key check and all: allowlists, limits, licences
left out, budgets and accounting apply as to any request, and a job is
the caller's. What a tool makes is kept as the caller's file and answered
by id and download URL, since a server over HTTP cannot write to the
caller's disk; a picture comes with a 512 px preview, and whole where no
file can be kept. A picture to work from is a file id, a data URL or
base64; one sent to the vision model is shrunk to 1280 px.

A mesh takes minutes. A client that takes a stream is answered as
server-sent events, with notifications/progress for its token while the
job is queued, waiting for the GPU and running, and waits up to ten
minutes; one that does not waits 50 seconds and gets the job id, which
job_status asks after.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-14 13:11:09 +02:00
9 changed files with 1314 additions and 10 deletions

No files matched your search

+12
View File
@@ -28,6 +28,18 @@ 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.
- **Extended thinking on `/v1/messages`**: `thinking: {type: "enabled"}`
tells a model that thinks by its chat template to (`enable_thinking`),
and its reasoning comes back as `thinking` blocks before the text,
+1 -1
View File
@@ -82,7 +82,7 @@
"/mcp": {
"post": {
"operationId": "postMcp",
"summary": "The gateway's MCP server: the memory notes as tools",
"summary": "The gateway's MCP server: the memory notes, and pictures, 3D models, vision and workflows, as tools",
"tags": [
"mcp"
],
+36
View File
@@ -534,6 +534,42 @@ 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
+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 as tools",
{Method: "POST", Path: "/mcp", Summary: "The gateway's MCP server: the memory notes, and pictures, 3D models, vision and workflows, 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{}},
+6 -2
View File
@@ -359,7 +359,9 @@ func (r *Router) Handler() http.Handler {
}
if r.auth == nil {
return r.refuseWhileStopped(r.refuseLoops(mux))
h := r.refuseWhileStopped(r.refuseLoops(mux))
r.selfHandler.Store(&h)
return h
}
exempt := []string{"/health", "/ready", "/leader", "/v1/openapi.json", "/", "/about", "/_static/", "/_lang/", "/claim/",
"/signup", "/login", "/login/", "/logout", "/profile", "/profile/", "/scim/",
@@ -373,7 +375,9 @@ 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.
return r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
h := r.refuseWhileStopped(r.refuseLoops(auth.Middleware(r.auth, r.log, exempt...)(mux)))
r.selfHandler.Store(&h)
return h
}
// refuseWhileStopped answers every request for a model with 503 while the
+937
View File
@@ -0,0 +1,937 @@
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
@@ -0,0 +1,266 @@
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)
}
}
+51 -6
View File
@@ -5,6 +5,7 @@ import (
"errors"
"io"
"net/http"
"slices"
"strconv"
"time"
@@ -120,7 +121,8 @@ func (r *Router) handlePutNote(w http.ResponseWriter, req *http.Request) {
Text string `json:"text"`
Tags []string `json:"tags"`
}
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
// Pictures to look at come in the call, as base64.
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, mcpBody))
if err != nil || json.Unmarshal(raw, &body) != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error",
`body must be JSON like {"text": "...", "tags": ["..."]}`)
@@ -273,23 +275,66 @@ func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
resp := rpcResponse{JSONRPC: "2.0", ID: rpc.ID}
switch rpc.Method {
case "initialize":
resp.Result = map[string]any{
"protocolVersion": mcpProtocolVer,
result := map[string]any{
"protocolVersion": mcpVersionFor(rpc.Params),
"capabilities": map[string]any{"tools": map[string]any{}},
"serverInfo": map[string]any{"name": "llmgw-memory", "version": "1"},
"serverInfo": map[string]any{"name": "llmgw", "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": mcpTools}
resp.Result = map[string]any{"tools": append(slices.Clone(mcpTools), r.mediaTools()...)}
case "tools/call":
resp.Result = r.callMemoryTool(req, rpc.Params)
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)
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,6 +51,10 @@ 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