gateway: a key held to some models counts no tokens and decides on none it may not use (#882, LCorleone) — counting tokens asked a provider's help or estimated locally whatever the key was held to, and System One asked its decision model of any key. Both now meet the hold the model endpoints do: count_tokens (Anthropic's shape) and Gemini's :countTokens — which the gateway answers itself, so no shared path covered it — resolve the model named and are refused with a 403 in the API's own error shape when the key may not use it, a group judged by its members, a bare group's name and an auto group's stand-in resolved first (countHeld in key_models.go); System One is held to the decision model its prefix names, after RouteDecider. TestGatewayKeyModelsHoldCountTokensAndSystemOne (403s by id, bare name and through a group on both counting APIs, Anthropic's and Gemini's error shapes, a held count or decision reaching no provider, an allowed model counted, a free key going on, the decider among the key's models) fails without the change. (#902)
This commit is contained in:
july_lu authored and GitHub committed
1 parent
3c7ceb68f4
commit
d7a8ddc84d
4 files changed
+144
No files matched your search
@@ -422,6 +422,12 @@ func (s *Server) serveSystemOne(w http.ResponseWriter, r *http.Request) {
|
||||
if asked == "" {
|
||||
asked = p.ID + "/" + model
|
||||
}
|
||||
// a gateway key held to some models (#882) is refused another; a
|
||||
// decision model is a model, and System One asks one directly
|
||||
if keyWho, held := keyHolds(r); held && !modelAllowed(keyWho, p, model) {
|
||||
writeError(w, provider.Chat, http.StatusForbidden, keyModelError(keyWho, asked))
|
||||
return
|
||||
}
|
||||
seat := decideSeat(p, model)
|
||||
var used Usage
|
||||
tr := s.trace.begin(Route{Time: start, Agent: agentOf(r), Session: sessionOf(r.Header), Model: asked, Provider: p.ID,
|
||||
|
||||
@@ -831,6 +831,11 @@ func (s *Server) countTokens(w http.ResponseWriter, r *http.Request) {
|
||||
id = sid
|
||||
}
|
||||
p, model, ok := provider.Resolve(id)
|
||||
// a gateway key held to some models (#882) counts no tokens on one it
|
||||
// may not use, as it would be refused serving it
|
||||
if countHeld(w, r, provider.Anthropic, id) {
|
||||
return
|
||||
}
|
||||
s.countOn(w, r, p, model, ok, body)
|
||||
}
|
||||
|
||||
@@ -992,6 +997,11 @@ func (s *Server) gemini(w http.ResponseWriter, r *http.Request) {
|
||||
case "streamGenerateContent":
|
||||
stream = true
|
||||
case "countTokens":
|
||||
// a gateway key held to some models (#882) counts no tokens on one
|
||||
// it may not use, as it would be refused serving it
|
||||
if countHeld(w, r, provider.Gemini, model) {
|
||||
return
|
||||
}
|
||||
s.geminiCount(w, model, body)
|
||||
return
|
||||
default:
|
||||
|
||||
@@ -135,3 +135,32 @@ func keyModelError(who access.Identity, model string) string {
|
||||
return fmt.Sprintf("The gateway key %q may not use %s; it may use %s. Change the key's models in magpie's Gateway page, or use a model it has.",
|
||||
who.KeyName, model, strings.Join(who.Models, ", "))
|
||||
}
|
||||
|
||||
// countHeld answers the 403 a gateway key held to some models (#882) gets
|
||||
// for counting tokens on one it may not use, as it would be refused
|
||||
// serving it: Anthropic's count_tokens and Gemini's :countTokens, id the
|
||||
// model as each asks it, resolved — a bare group's name, an auto group's
|
||||
// stand-in — and a group judged by its members. proto is the API the
|
||||
// error is answered on.
|
||||
func countHeld(w http.ResponseWriter, r *http.Request, proto provider.Protocol, id string) bool {
|
||||
keyWho, held := keyHolds(r)
|
||||
if !held {
|
||||
return false
|
||||
}
|
||||
if gid, ok := provider.GroupFor(id); ok {
|
||||
id = gid
|
||||
}
|
||||
if sid, ok := provider.AutoStandIn(id); ok {
|
||||
id = sid
|
||||
}
|
||||
p, model, ok := provider.Resolve(id)
|
||||
if !ok {
|
||||
return false // counted nowhere: a local estimate, as before
|
||||
}
|
||||
g, ms, isGroup := provider.FindGroup(id)
|
||||
if isGroup && !groupAllowed(keyWho, g, ms) || !isGroup && !modelAllowed(keyWho, p, model) {
|
||||
writeError(w, proto, http.StatusForbidden, keyModelError(keyWho, id))
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -2,6 +2,8 @@ package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"strings"
|
||||
@@ -242,3 +244,100 @@ func TestGatewayKeyModelsNameAGroup(t *testing.T) {
|
||||
t.Fatal("a named group fell back outside itself", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// A gateway key held to some models (#882) counts no tokens on one it may
|
||||
// not use — Anthropic's count_tokens and Gemini's :countTokens, a group
|
||||
// judged by its members — and System One asks its decision model only of
|
||||
// a key that may use it; a free key goes on as before, and no provider is
|
||||
// asked for any of it.
|
||||
func TestGatewayKeyModelsHoldCountTokensAndSystemOne(t *testing.T) {
|
||||
fresh(t)
|
||||
plan := &fake{t: t, ctype: "application/json", reply: `{"id":"from-plan","choices":[]}`}
|
||||
spare := &fake{t: t, ctype: "application/json", reply: `{"id":"from-spare","choices":[]}`}
|
||||
twoProviders(t, plan, spare)
|
||||
var decided int
|
||||
jev := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
decided++
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
io.WriteString(w, `{"answers":{"intent":{"choice":"bug"}},"usage":{"input_tokens":10,"output_tokens":2}}`)
|
||||
}))
|
||||
t.Cleanup(jev.Close)
|
||||
if err := provider.Save(provider.Provider{ID: "jev", Name: "Jev", Key: "k", Decide: jev.URL + "/v1"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := provider.SaveGroup(provider.Group{Name: "Both", Members: []string{"plan/m1", "spare/m2"}, Routing: provider.Ordered}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys, secrets := newCaller(t, "Held", "Free")
|
||||
if _, err := access.Update("models-key", access.Change{Key: keys[0].ID, Models: []string{"plan/*"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := New().Handler()
|
||||
do := func(secret, method, path, body string) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
r.Header.Set("Authorization", "Bearer "+secret)
|
||||
w := httptest.NewRecorder()
|
||||
h.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
// counting tokens on a model it may not use, by id, bare and through a
|
||||
// group with one in: said in Anthropic's own error shape
|
||||
for _, model := range []string{"spare/m2", "m2", "group/both"} {
|
||||
w := do(secrets[0], "POST", "/v1/messages/count_tokens", `{"model":"`+model+`","messages":[{"role":"user","content":"hi"}]}`)
|
||||
var e struct {
|
||||
Error struct{ Message, Type string }
|
||||
}
|
||||
json.Unmarshal(w.Body.Bytes(), &e)
|
||||
if w.Code != 403 || e.Error.Type != "permission_error" || !strings.Contains(e.Error.Message, `"Held"`) || !strings.Contains(e.Error.Message, "plan/*") {
|
||||
t.Fatalf("counting %s: %d %s", model, w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
// one it may use is counted, and a free key counts anything
|
||||
for _, c := range []struct{ secret, model string }{
|
||||
{secrets[0], "plan/m1"}, {secrets[1], "spare/m2"},
|
||||
} {
|
||||
if w := do(c.secret, "POST", "/v1/messages/count_tokens", `{"model":"`+c.model+`","messages":[{"role":"user","content":"hi"}]}`); w.Code != 200 || !strings.Contains(w.Body.String(), "input_tokens") {
|
||||
t.Fatalf("counting %s: %d %s", c.model, w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
// Gemini's counting, the model in its URL, is held the same way
|
||||
w := do(secrets[0], "POST", "/v1beta/models/spare/m2:countTokens", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)
|
||||
var g struct {
|
||||
Error struct {
|
||||
Code int
|
||||
Status string
|
||||
}
|
||||
}
|
||||
if w.Code != 403 || json.Unmarshal(w.Body.Bytes(), &g) != nil || g.Error.Code != 403 || g.Error.Status != "PERMISSION_DENIED" {
|
||||
t.Fatalf("gemini counting: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
if w = do(secrets[0], "POST", "/v1beta/models/plan/m1:countTokens", `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`); w.Code != 200 || !strings.Contains(w.Body.String(), "totalTokens") {
|
||||
t.Fatalf("gemini counting allowed: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
if plan.calls != 0 || spare.calls != 0 || decided != 0 {
|
||||
t.Fatal("a held count or decision reached a provider", plan.calls, spare.calls, decided)
|
||||
}
|
||||
// System One's decision model is a model: the held key is refused it,
|
||||
// a free key asks, and one given the decider asks too
|
||||
const ask = `{"model":"jev/jev-latest","state":{"message":"hi"},"questions":{"intent":{"type":"choice"}}}`
|
||||
w = do(secrets[0], "POST", "/v1/systemone", ask)
|
||||
var e struct {
|
||||
Error struct{ Message, Type string }
|
||||
}
|
||||
json.Unmarshal(w.Body.Bytes(), &e)
|
||||
if w.Code != 403 || e.Error.Type != "permission_error" || !strings.Contains(e.Error.Message, `"Held"`) {
|
||||
t.Fatalf("system one held: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
if decided != 0 {
|
||||
t.Fatal("a held decision reached its provider", decided)
|
||||
}
|
||||
if w = do(secrets[1], "POST", "/v1/systemone", ask); w.Code != 200 || decided != 1 {
|
||||
t.Fatalf("a free key's decision: %d %s, decided %d", w.Code, w.Body.String(), decided)
|
||||
}
|
||||
if _, err := access.Update("models-key", access.Change{Key: keys[0].ID, Models: []string{"jev/*"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w = do(secrets[0], "POST", "/v1/systemone", ask); w.Code != 200 || decided != 2 {
|
||||
t.Fatalf("a key given the decider: %d %s, decided %d", w.Code, w.Body.String(), decided)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user