From 55d5e13534b4946b5f049d904570af513a20b67f Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Wed, 14 Oct 2026 13:11:09 +0200 Subject: [PATCH] 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 --- CHANGELOG.md | 12 + docs/client-api.openapi.json | 2 +- docs/clients.md | 36 ++ internal/clientapi/routes.go | 2 +- internal/router/handlers.go | 8 +- internal/router/mcp_media.go | 937 ++++++++++++++++++++++++++++++ internal/router/mcp_media_test.go | 266 +++++++++ internal/router/memory.go | 57 +- internal/router/router.go | 4 + 9 files changed, 1314 insertions(+), 10 deletions(-) create mode 100644 internal/router/mcp_media.go create mode 100644 internal/router/mcp_media_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index c9821cb..fa69aeb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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, diff --git a/docs/client-api.openapi.json b/docs/client-api.openapi.json index 9548e3c..048ba00 100644 --- a/docs/client-api.openapi.json +++ b/docs/client-api.openapi.json @@ -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" ], diff --git a/docs/clients.md b/docs/clients.md index 982f765..9936ef3 100644 --- a/docs/clients.md +++ b/docs/clients.md @@ -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 `). 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 diff --git a/internal/clientapi/routes.go b/internal/clientapi/routes.go index b1f24ef..8c87b82 100644 --- a/internal/clientapi/routes.go +++ b/internal/clientapi/routes.go @@ -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{}}, diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 741d396..ffe01a7 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -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 diff --git a/internal/router/mcp_media.go b/internal/router/mcp_media.go new file mode 100644 index 0000000..70b3ad7 --- /dev/null +++ b/internal/router/mcp_media.go @@ -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") +} diff --git a/internal/router/mcp_media_test.go b/internal/router/mcp_media_test.go new file mode 100644 index 0000000..235dd64 --- /dev/null +++ b/internal/router/mcp_media_test.go @@ -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) + } +} diff --git a/internal/router/memory.go b/internal/router/memory.go index 495e63b..d5c98f0 100644 --- a/internal/router/memory.go +++ b/internal/router/memory.go @@ -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}}, diff --git a/internal/router/router.go b/internal/router/router.go index cfea90c..136bc6d 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -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