chore: initial commit
This commit is contained in:
@@ -0,0 +1,76 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/netip"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Server struct {
|
||||
TrustedProxies []string `yaml:"trusted_proxies"`
|
||||
Listen string `yaml:"listen"`
|
||||
PublicOrigin string `yaml:"public_origin"`
|
||||
} `yaml:"server"`
|
||||
Admin struct {
|
||||
Password string `yaml:"password"`
|
||||
} `yaml:"admin"`
|
||||
Storage struct {
|
||||
LogRetentionDays int `yaml:"log_retention_days"`
|
||||
Path string `yaml:"path"`
|
||||
} `yaml:"storage"`
|
||||
Limits struct {
|
||||
RequestsPerMinute int `yaml:"requests_per_minute"`
|
||||
PerKeyConcurrency int `yaml:"per_key_concurrency"`
|
||||
GlobalConcurrency int `yaml:"global_concurrency"`
|
||||
} `yaml:"limits"`
|
||||
}
|
||||
|
||||
func Load(path string) (Config, error) {
|
||||
var c Config
|
||||
c.Limits.RequestsPerMinute = 120
|
||||
c.Limits.PerKeyConcurrency = 8
|
||||
c.Limits.GlobalConcurrency = 64
|
||||
c.Server.Listen = "127.0.0.1:8080"
|
||||
c.Storage.Path = "./data/router.db"
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return c, err
|
||||
}
|
||||
defer f.Close()
|
||||
dec := yaml.NewDecoder(f)
|
||||
dec.KnownFields(true)
|
||||
if err = dec.Decode(&c); err != nil {
|
||||
return c, fmt.Errorf("invalid config: %w", err)
|
||||
}
|
||||
var extra any
|
||||
if err = dec.Decode(&extra); err != io.EOF {
|
||||
return c, fmt.Errorf("config must contain exactly one YAML document")
|
||||
}
|
||||
if len(c.Admin.Password) < 12 || strings.HasPrefix(c.Admin.Password, "CHANGE-ME") {
|
||||
return c, fmt.Errorf("admin.password must be a non-placeholder password of at least 12 characters")
|
||||
}
|
||||
if c.Server.Listen == "" || c.Storage.Path == "" {
|
||||
return c, fmt.Errorf("server.listen and storage.path must not be empty")
|
||||
}
|
||||
if c.Limits.RequestsPerMinute <= 0 || c.Limits.PerKeyConcurrency <= 0 || c.Limits.GlobalConcurrency <= 0 {
|
||||
return c, fmt.Errorf("limits must be positive integers")
|
||||
}
|
||||
if c.Storage.LogRetentionDays < 0 {
|
||||
return c, fmt.Errorf("storage.log_retention_days must be nonnegative")
|
||||
}
|
||||
for _, cidr := range c.Server.TrustedProxies {
|
||||
p, err := netip.ParsePrefix(cidr)
|
||||
if err != nil {
|
||||
return c, fmt.Errorf("invalid trusted proxy CIDR %q: %w", cidr, err)
|
||||
}
|
||||
if p.Bits() == 0 {
|
||||
return c, fmt.Errorf("trusted proxy CIDR must not trust all addresses")
|
||||
}
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoad(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
name, body string
|
||||
wantErr bool
|
||||
}{
|
||||
{"valid", "admin:\n password: a-long-test-password\n", false},
|
||||
{"short", "admin:\n password: short\n", true},
|
||||
{"multiple documents", "admin:\n password: a-long-test-password\n---\nadmin:\n password: ignored-password\n", true},
|
||||
{"placeholder", "admin:\n password: CHANGE-ME-use-a-long-random-password\n", true},
|
||||
{"unknown", "admin:\n password: a-long-test-password\nunknown: true\n", true},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.yaml")
|
||||
if err := os.WriteFile(path, []byte(tt.body), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := Load(path)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Fatalf("error=%v, wantErr=%v", err, tt.wantErr)
|
||||
}
|
||||
if err == nil && (c.Server.Listen != "127.0.0.1:8080" || c.Storage.Path != "./data/router.db") {
|
||||
t.Fatal("defaults missing")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityConfig(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
body string
|
||||
invalid bool
|
||||
}{
|
||||
{"", false},
|
||||
{"limits:\n requests_per_minute: 0\n", true},
|
||||
{"limits:\n per_key_concurrency: -1\n", true},
|
||||
{"limits:\n global_concurrency: 0\n", true},
|
||||
{"limits:\n requests_per_minute: nope\n", true},
|
||||
{"server:\n trusted_proxies: [not-a-cidr]\n", true},
|
||||
{"server:\n trusted_proxies: [0.0.0.0/0]\n", true},
|
||||
{"server:\n trusted_proxies: ['::/0']\n", true},
|
||||
{"server:\n trusted_proxies: [10.0.0.0/8, '::1/128']\n", false},
|
||||
{"storage:\n log_retention_days: -1\n", true},
|
||||
} {
|
||||
t.Run(tt.body, func(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "config.yaml")
|
||||
if err := os.WriteFile(p, []byte("admin:\n password: a-long-test-password\n"+tt.body), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := Load(p)
|
||||
if (err != nil) != tt.invalid {
|
||||
t.Fatalf("error=%v", err)
|
||||
}
|
||||
if err == nil && (c.Limits.RequestsPerMinute != 120 || c.Limits.PerKeyConcurrency != 8 || c.Limits.GlobalConcurrency != 64 || c.Storage.LogRetentionDays != 0) {
|
||||
t.Fatal("security defaults missing")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
# Protocol conversion
|
||||
|
||||
Public entry points are `ConvertRequest`, `ConvertResponse`, `Stream`, and
|
||||
`ExtractUsage`. No external dependencies are needed. The service should bypass
|
||||
conversion for same-protocol non-streaming traffic. `Stream` preserves original
|
||||
SSE bytes for same-protocol traffic, including comments and CRLF, and returns
|
||||
upstream errors instead of silently swallowing them.
|
||||
|
||||
Cross-protocol support includes text, user URL/base64 images, system messages,
|
||||
function schemas and selection, multiple function calls, tool results, multi-turn
|
||||
history, and incremental tool argument JSON. Responses output item/content part/
|
||||
argument start, delta, done and terminal events are emitted. Anthropic content
|
||||
block and message lifecycles are emitted. Each write flushes when supported.
|
||||
|
||||
## Deliberate limits
|
||||
|
||||
- Anthropic output thinking/redacted_thinking blocks and their thinking/signature
|
||||
stream deltas are omitted when converting to Chat/Responses, never appended to
|
||||
visible text. Block indices and lifecycle are checked. Native Anthropic output
|
||||
bytes remain unchanged; request-side thinking is not translated.
|
||||
|
||||
- Chat output message/delta reasoning_content (string or null) is omitted only
|
||||
during cross-protocol conversion; it is never appended to visible text or
|
||||
fabricated as Anthropic thinking. Native Chat responses/SSE remain unchanged.
|
||||
Optional indices on complete Chat tool calls must match their array positions.
|
||||
|
||||
- Responses output message `phase` is accepted as metadata. Cross-protocol
|
||||
conversion deliberately omits Responses output reasoning, summaries and
|
||||
encrypted content, including their recognized stream events. They are never
|
||||
appended to visible text or converted to Anthropic thinking/signatures.
|
||||
Same-protocol responses and SSE retain these bytes unchanged. Reasoning
|
||||
item/part indices, identities and lifecycles are checked; unrelated unknown
|
||||
content and event types remain errors. Empty output-text logprobs arrays
|
||||
are accepted, but actual logprobs remain unsupported.
|
||||
- Only one choice and function tools are supported. Other reasoning/thinking blocks,
|
||||
audio/video/documents, hosted tools, citations, refusal content, logprobs,
|
||||
structured-output modes, prompt-cache controls and Responses server-side
|
||||
conversation state are rejected rather than approximated. Unknown request
|
||||
fields and unknown stream event types are rejected.
|
||||
- Anthropic targets require an explicit token limit. Developer roles,
|
||||
mid-conversation system messages, strict tool schemas and explicit image
|
||||
detail levels cannot be translated to Anthropic. The implementation currently
|
||||
conservatively rejects explicit image detail for all cross-protocol pairs.
|
||||
- Tool results must be text and cannot carry Anthropic's `is_error=true`.
|
||||
Stop sequences cannot target Responses; stop-sequence termination cannot
|
||||
be represented losslessly across protocols.
|
||||
- Responses streaming supports one output-text content part per message.
|
||||
Final content must agree with deltas; nonincremental/final-only variants
|
||||
are rejected. Tool identity must be complete before argument streaming;
|
||||
later ID/name fragments are rejected. Parallel tool argument fragments are
|
||||
tracked independently.
|
||||
- `Usage.Input` includes cache-read and cache-write tokens. Anthropic input
|
||||
counters exclude these, so they are added on extraction and subtracted on
|
||||
encoding. `Usage.Raw` holds upstream usage JSON; Anthropic stream snapshots
|
||||
are merged so final output-only usage does not erase initial input counts.
|
||||
Unknown usage remains `Known=false`; Anthropic wire messages still carry
|
||||
required zero counters until actual usage is available. OpenAI wire schemas
|
||||
have no separate Anthropic cache-write counter; it remains in returned Usage.
|
||||
- SSE events are limited to 16 MiB; generated content is retained for final
|
||||
Responses snapshots and argument validation. LF and CRLF framing is supported
|
||||
(not standalone CR). Same-protocol streams allow EOF without a terminal marker;
|
||||
cross-protocol streams report truncation. Terminal markers finish immediately.
|
||||
- Context is checked between reads and writes. A generic `io.Reader` cannot be
|
||||
interrupted safely: the service must bind upstream HTTP requests to the same
|
||||
context or close their body on cancellation. Writers/readers remain caller-owned.
|
||||
- Response transport metadata and echoed Responses request settings are not
|
||||
round-tripped. Generated IDs use cryptographic randomness if upstream omits
|
||||
an ID. No external provider or SDK integration tests have been performed.
|
||||
|
||||
Run `go test ./internal/protocol`, `go test -race -cover ./internal/protocol`,
|
||||
and `go vet ./internal/protocol`.
|
||||
@@ -0,0 +1,101 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAnthropicThinkingOutput(t *testing.T) {
|
||||
for _, tool := range []bool{false, true} {
|
||||
visible := `{"type":"text","text":"OK"}`
|
||||
initial := `{"type":"text","text":""}`
|
||||
delta := `{"type":"text_delta","text":"OK"}`
|
||||
finish := "end_turn"
|
||||
if tool {
|
||||
visible = `{"type":"tool_use","id":"call_1","name":"weather","input":{}}`
|
||||
initial = visible
|
||||
delta = `{"type":"input_json_delta","partial_json":"{}"}`
|
||||
finish = "tool_use"
|
||||
}
|
||||
hidden := `{"type":"thinking","thinking":"PRIVATE","signature":"SIGNATURE"}`
|
||||
redacted := `{"type":"redacted_thinking","data":"CIPHER"}`
|
||||
body := `{"id":"m","type":"message","role":"assistant","model":"m","content":[` + hidden + "," + redacted + "," + visible + `],"stop_reason":"` + finish + `","usage":{"input_tokens":9,"output_tokens":7}}`
|
||||
events := []string{
|
||||
`{"type":"message_start","message":{"id":"m","content":[],"usage":{"input_tokens":9,"output_tokens":0}}}`,
|
||||
`{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}`,
|
||||
`{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"PRIVATE"}}`,
|
||||
`{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"SIGNATURE"}}`,
|
||||
`{"type":"content_block_stop","index":0}`,
|
||||
`{"type":"content_block_start","index":1,"content_block":` + redacted + "}",
|
||||
`{"type":"content_block_stop","index":1}`,
|
||||
`{"type":"content_block_start","index":2,"content_block":` + initial + "}",
|
||||
`{"type":"content_block_delta","index":2,"delta":` + delta + "}",
|
||||
`{"type":"content_block_stop","index":2}`,
|
||||
`{"type":"message_delta","delta":{"stop_reason":"` + finish + `"},"usage":{"output_tokens":7}}`,
|
||||
`{"type":"message_stop"}`,
|
||||
}
|
||||
src := reasoningSSE(events...)
|
||||
for _, target := range []Kind{Chat, Responses, Anthropic} {
|
||||
out, err := ConvertResponse(Anthropic, target, []byte(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var dst bytes.Buffer
|
||||
u, err := Stream(context.Background(), &dst, strings.NewReader(src), Anthropic, target, "m")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !u.Known || u.Input != 9 || u.Output != 7 {
|
||||
t.Fatalf("usage %+v", u)
|
||||
}
|
||||
if target == Anthropic {
|
||||
if string(out) != body || dst.String() != src {
|
||||
t.Fatal("native bytes changed")
|
||||
}
|
||||
continue
|
||||
}
|
||||
for _, result := range []string{string(out), dst.String()} {
|
||||
for _, secret := range []string{"PRIVATE", "SIGNATURE", "CIPHER", "thinking"} {
|
||||
if strings.Contains(result, secret) {
|
||||
t.Fatal("thinking leaked")
|
||||
}
|
||||
}
|
||||
want := "OK"
|
||||
if tool {
|
||||
want = "weather"
|
||||
}
|
||||
if !strings.Contains(result, want) {
|
||||
t.Fatal("lost visible output")
|
||||
}
|
||||
}
|
||||
terminal := "[DONE]"
|
||||
if target == Responses {
|
||||
terminal = "response.completed"
|
||||
}
|
||||
if !strings.Contains(dst.String(), terminal) {
|
||||
t.Fatal("no terminal")
|
||||
}
|
||||
}
|
||||
for _, bad := range []string{
|
||||
strings.Replace(src, `{"type":"content_block_stop","index":0}`, `{"type":"ping"}`, 1),
|
||||
strings.Replace(src, `"type":"thinking_delta"`, `"type":"text_delta"`, 1),
|
||||
strings.Replace(src, `"index":2,"content_block"`, `"index":0,"content_block"`, 1),
|
||||
} {
|
||||
var dst bytes.Buffer
|
||||
if _, err := Stream(context.Background(), &dst, strings.NewReader(bad), Anthropic, Chat, "m"); err == nil {
|
||||
t.Fatal("accepted invalid thinking lifecycle")
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, block := range []object{
|
||||
{"type": "thinking", "thinking": true, "signature": ""},
|
||||
{"type": "thinking", "thinking": "", "signature": "", "unknown": true},
|
||||
{"type": "redacted_thinking", "data": nil},
|
||||
} {
|
||||
if _, err := anthropicThinking(block); err == nil {
|
||||
t.Fatal("accepted malformed thinking")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestChatOutputReasoning(t *testing.T) {
|
||||
for _, tool := range []bool{false, true} {
|
||||
message := object{"role": "assistant", "content": "OK", "reasoning_content": "PRIVATE"}
|
||||
delta := object{"content": "OK", "reasoning_content": "PRIVATE"}
|
||||
finish := "stop"
|
||||
if tool {
|
||||
call := object{"index": 0, "id": "call_1", "type": "function", "function": object{"name": "weather", "arguments": "{}"}}
|
||||
message["content"] = nil
|
||||
message["tool_calls"] = []any{call}
|
||||
delta = object{"tool_calls": []any{call}, "reasoning_content": nil}
|
||||
finish = "tool_calls"
|
||||
}
|
||||
usage := object{"prompt_tokens": 9, "completion_tokens": 7}
|
||||
response := object{"id": "c", "object": "chat.completion", "model": "m", "choices": []any{object{"index": 0, "message": message, "finish_reason": finish}}, "usage": usage}
|
||||
body, _ := json.Marshal(response)
|
||||
chunk, _ := json.Marshal(object{"id": "c", "choices": []any{object{"index": 0, "delta": delta, "finish_reason": finish}}, "usage": usage})
|
||||
reasoning, _ := json.Marshal(object{"id": "c", "choices": []any{object{"index": 0, "delta": object{"role": "assistant", "reasoning_content": "PRIVATE"}, "finish_reason": nil}}})
|
||||
sep := string([]byte{10, 10})
|
||||
src := "data: " + string(reasoning) + sep + "data: " + string(chunk) + sep + "data: [DONE]" + sep
|
||||
for _, target := range []Kind{Chat, Responses, Anthropic} {
|
||||
out, err := ConvertResponse(Chat, target, body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var dst bytes.Buffer
|
||||
u, err := Stream(context.Background(), &dst, strings.NewReader(src), Chat, target, "m")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !u.Known || u.Input != 9 || u.Output != 7 {
|
||||
t.Fatalf("usage: %+v", u)
|
||||
}
|
||||
if target == Chat {
|
||||
if !bytes.Equal(out, body) || dst.String() != src {
|
||||
t.Fatal("native bytes changed")
|
||||
}
|
||||
continue
|
||||
}
|
||||
for _, text := range []string{string(out), dst.String()} {
|
||||
if strings.Contains(text, "PRIVATE") || strings.Contains(text, "reasoning_content") || strings.Contains(text, "thinking") {
|
||||
t.Fatal("reasoning leaked")
|
||||
}
|
||||
want := "OK"
|
||||
if tool {
|
||||
want = "weather"
|
||||
}
|
||||
if !strings.Contains(text, want) {
|
||||
t.Fatal("lost output")
|
||||
}
|
||||
}
|
||||
terminal := "response.completed"
|
||||
if target == Anthropic {
|
||||
terminal = "message_stop"
|
||||
}
|
||||
if !strings.Contains(dst.String(), terminal) {
|
||||
t.Fatal("no terminal")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatOutputMetadataValidation(t *testing.T) {
|
||||
for _, val := range []any{nil, "", "private"} {
|
||||
if _, err := chatOutput(object{"reasoning_content": val}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
for _, val := range []any{true, object{}, []any{}} {
|
||||
if _, err := chatOutput(object{"reasoning_content": val}); err == nil {
|
||||
t.Fatal("accepted malformed reasoning")
|
||||
}
|
||||
}
|
||||
for _, index := range []any{json.Number("1"), json.Number("-1"), json.Number("0.5"), "0", nil} {
|
||||
if _, err := chatResponseMessage(object{"tool_calls": []any{object{"index": index}}}); err == nil {
|
||||
t.Fatal("accepted malformed index")
|
||||
}
|
||||
}
|
||||
body := []byte(`{"id":"c","choices":[{"message":{"role":"assistant","content":"OK","reasoning_content":"hidden","unknown_field":true},"finish_reason":"stop"}]}`)
|
||||
if _, err := ConvertResponse(Chat, Responses, body); err == nil {
|
||||
t.Fatal("accepted unknown field")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,516 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newID() string { return fmt.Sprintf("resp_%x", randomID()) }
|
||||
func randomID() []byte {
|
||||
b := make([]byte, 16)
|
||||
if _, e := rand.Read(b); e != nil {
|
||||
panic(e)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func contentBlock(k Kind, p part, role string) (object, error) {
|
||||
switch p.typ {
|
||||
case "text":
|
||||
typ := "text"
|
||||
if k == Responses {
|
||||
typ = "input_text"
|
||||
if role == "assistant" {
|
||||
typ = "output_text"
|
||||
}
|
||||
}
|
||||
return object{"type": typ, "text": p.text}, nil
|
||||
case "image":
|
||||
if role != "user" {
|
||||
return nil, fmt.Errorf("images only supported in user messages")
|
||||
}
|
||||
switch k {
|
||||
case Chat:
|
||||
return object{"type": "image_url", "image_url": object{"url": p.url}}, nil
|
||||
case Responses:
|
||||
return object{"type": "input_image", "image_url": p.url}, nil
|
||||
default:
|
||||
src := object{"type": "url", "url": p.url}
|
||||
if strings.HasPrefix(p.url, "data:") {
|
||||
a, b, ok := strings.Cut(strings.TrimPrefix(p.url, "data:"), ";base64,")
|
||||
if !ok || a == "" || b == "" {
|
||||
return nil, fmt.Errorf("unsupported image data URL")
|
||||
}
|
||||
src = object{"type": "base64", "media_type": a, "data": b}
|
||||
}
|
||||
return object{"type": "image", "source": src}, nil
|
||||
}
|
||||
case "call":
|
||||
if role != "assistant" {
|
||||
return nil, fmt.Errorf("tool call outside assistant message")
|
||||
}
|
||||
switch k {
|
||||
case Chat:
|
||||
return object{"id": p.id, "type": "function", "function": object{"name": p.name, "arguments": p.args}}, nil
|
||||
case Responses:
|
||||
return object{"type": "function_call", "call_id": p.id, "name": p.name, "arguments": p.args}, nil
|
||||
default:
|
||||
args, e := decode([]byte(p.args))
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
return object{"type": "tool_use", "id": p.id, "name": p.name, "input": args}, nil
|
||||
}
|
||||
case "result":
|
||||
if role != "tool" && role != "user" {
|
||||
return nil, fmt.Errorf("tool result outside user/tool message")
|
||||
}
|
||||
switch k {
|
||||
case Chat:
|
||||
return object{"role": "tool", "tool_call_id": p.id, "content": p.text}, nil
|
||||
case Responses:
|
||||
return object{"type": "function_call_output", "call_id": p.id, "output": p.text}, nil
|
||||
default:
|
||||
return object{"type": "tool_result", "tool_use_id": p.id, "content": p.text}, nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("unsupported part")
|
||||
}
|
||||
func writeRequest(k Kind, r request) (object, error) {
|
||||
out := object{}
|
||||
for key, val := range r.options {
|
||||
switch key {
|
||||
case "max":
|
||||
switch k {
|
||||
case Chat:
|
||||
key = "max_completion_tokens"
|
||||
case Responses:
|
||||
key = "max_output_tokens"
|
||||
case Anthropic:
|
||||
key = "max_tokens"
|
||||
}
|
||||
case "stop":
|
||||
if k == Responses {
|
||||
return nil, fmt.Errorf("Responses does not support stop sequences")
|
||||
}
|
||||
if k == Anthropic {
|
||||
key = "stop_sequences"
|
||||
}
|
||||
case "parallel_tool_calls":
|
||||
if k == Anthropic {
|
||||
continue
|
||||
}
|
||||
}
|
||||
out[key] = val
|
||||
}
|
||||
if k == Anthropic {
|
||||
if _, ok := out["max_tokens"]; !ok {
|
||||
return nil, fmt.Errorf("Anthropic requires explicit max_tokens")
|
||||
}
|
||||
}
|
||||
ms := []any{}
|
||||
sys := []any{}
|
||||
seenNonSystem := false
|
||||
for _, m := range r.messages {
|
||||
if k == Anthropic && (m.role == "system" || m.role == "developer") {
|
||||
if m.role == "developer" {
|
||||
return nil, fmt.Errorf("developer role cannot be represented by Anthropic")
|
||||
}
|
||||
if seenNonSystem {
|
||||
return nil, fmt.Errorf("mid-conversation system message cannot be represented by Anthropic")
|
||||
}
|
||||
for _, p := range m.parts {
|
||||
if p.typ != "text" {
|
||||
return nil, fmt.Errorf("non-text system content")
|
||||
}
|
||||
sys = append(sys, object{"type": "text", "text": p.text})
|
||||
}
|
||||
continue
|
||||
}
|
||||
seenNonSystem = true
|
||||
role := m.role
|
||||
if k == Anthropic && role == "tool" {
|
||||
role = "user"
|
||||
}
|
||||
cs := []any{}
|
||||
calls := []any{}
|
||||
flush := func() {
|
||||
if len(cs) > 0 {
|
||||
ms = append(ms, object{"role": role, "content": cs})
|
||||
cs = []any{}
|
||||
}
|
||||
}
|
||||
for _, p := range m.parts {
|
||||
b, e := contentBlock(k, p, m.role)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
switch {
|
||||
case k == Chat && p.typ == "call":
|
||||
calls = append(calls, b)
|
||||
case (k == Chat && p.typ == "result") || (k == Responses && (p.typ == "call" || p.typ == "result")):
|
||||
flush()
|
||||
ms = append(ms, b)
|
||||
default:
|
||||
cs = append(cs, b)
|
||||
}
|
||||
}
|
||||
if k == Chat && len(calls) > 0 {
|
||||
x := object{"role": role, "tool_calls": calls, "content": nil}
|
||||
if len(cs) > 0 {
|
||||
x["content"] = cs
|
||||
}
|
||||
ms = append(ms, x)
|
||||
} else {
|
||||
flush()
|
||||
}
|
||||
}
|
||||
if k == Anthropic {
|
||||
merged := []any{}
|
||||
for _, v := range ms {
|
||||
m := obj(v)
|
||||
if len(merged) > 0 && obj(merged[len(merged)-1])["role"] == m["role"] {
|
||||
last := obj(merged[len(merged)-1])
|
||||
last["content"] = append(arr(last["content"]), arr(m["content"])...)
|
||||
} else {
|
||||
merged = append(merged, m)
|
||||
}
|
||||
}
|
||||
ms = merged
|
||||
if len(sys) > 0 {
|
||||
out["system"] = sys
|
||||
}
|
||||
}
|
||||
if k == Responses {
|
||||
out["input"] = ms
|
||||
out["store"] = false
|
||||
} else {
|
||||
out["messages"] = ms
|
||||
}
|
||||
if len(r.tools) > 0 {
|
||||
ts := []any{}
|
||||
for _, f := range r.tools {
|
||||
switch k {
|
||||
case Chat:
|
||||
ts = append(ts, object{"type": "function", "function": f})
|
||||
case Responses:
|
||||
t := object{"type": "function"}
|
||||
for key, v := range f {
|
||||
t[key] = v
|
||||
}
|
||||
ts = append(ts, t)
|
||||
case Anthropic:
|
||||
t := object{}
|
||||
for key, v := range f {
|
||||
if key == "strict" {
|
||||
if v != false {
|
||||
return nil, fmt.Errorf("strict tool schema cannot be enforced by Anthropic")
|
||||
}
|
||||
continue
|
||||
}
|
||||
if key == "parameters" {
|
||||
key = "input_schema"
|
||||
}
|
||||
t[key] = v
|
||||
}
|
||||
ts = append(ts, t)
|
||||
}
|
||||
}
|
||||
out["tools"] = ts
|
||||
}
|
||||
if r.choice != nil {
|
||||
if s, ok := r.choice.(string); ok {
|
||||
if k == Anthropic {
|
||||
if s == "required" {
|
||||
s = "any"
|
||||
}
|
||||
out["tool_choice"] = object{"type": s}
|
||||
} else {
|
||||
out["tool_choice"] = s
|
||||
}
|
||||
} else {
|
||||
name := obj(r.choice)["name"]
|
||||
switch k {
|
||||
case Chat:
|
||||
out["tool_choice"] = object{"type": "function", "function": object{"name": name}}
|
||||
case Responses:
|
||||
out["tool_choice"] = object{"type": "function", "name": name}
|
||||
case Anthropic:
|
||||
out["tool_choice"] = object{"type": "tool", "name": name}
|
||||
}
|
||||
}
|
||||
}
|
||||
if k == Anthropic {
|
||||
if v, ok := r.options["parallel_tool_calls"]; ok {
|
||||
b, ok := v.(bool)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid parallel_tool_calls")
|
||||
}
|
||||
tc := obj(out["tool_choice"])
|
||||
if tc == nil {
|
||||
tc = object{"type": "auto"}
|
||||
out["tool_choice"] = tc
|
||||
}
|
||||
tc["disable_parallel_tool_use"] = !b
|
||||
}
|
||||
}
|
||||
if k == Chat && out["stream"] == true {
|
||||
out["stream_options"] = object{"include_usage": true}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
type response struct {
|
||||
id, model, finish string
|
||||
parts []part
|
||||
usage Usage
|
||||
}
|
||||
|
||||
func finishReason(k Kind, s string) (string, error) {
|
||||
switch k {
|
||||
case Chat:
|
||||
switch s {
|
||||
case "stop":
|
||||
return "stop", nil
|
||||
case "tool_calls":
|
||||
return "tool", nil
|
||||
case "length":
|
||||
return "length", nil
|
||||
}
|
||||
case Anthropic:
|
||||
switch s {
|
||||
case "end_turn":
|
||||
return "stop", nil
|
||||
case "tool_use":
|
||||
return "tool", nil
|
||||
case "max_tokens":
|
||||
return "length", nil
|
||||
}
|
||||
case Responses:
|
||||
if s == "completed" {
|
||||
return "stop", nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("unsupported finish reason %q", s)
|
||||
}
|
||||
func wireFinish(k Kind, s string) string {
|
||||
switch k {
|
||||
case Chat:
|
||||
switch s {
|
||||
case "tool":
|
||||
return "tool_calls"
|
||||
case "length":
|
||||
return "length"
|
||||
}
|
||||
return "stop"
|
||||
case Anthropic:
|
||||
switch s {
|
||||
case "tool":
|
||||
return "tool_use"
|
||||
case "length":
|
||||
return "max_tokens"
|
||||
}
|
||||
return "end_turn"
|
||||
default:
|
||||
if s == "length" {
|
||||
return "incomplete"
|
||||
}
|
||||
return "completed"
|
||||
}
|
||||
}
|
||||
func ConvertResponse(from, to Kind, body []byte) ([]byte, error) {
|
||||
if !valid(from) || !valid(to) {
|
||||
return nil, fmt.Errorf("unknown protocol")
|
||||
}
|
||||
m, e := decode(body)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if from == to {
|
||||
return body, nil
|
||||
}
|
||||
r, e := readResponse(from, m)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
out, e := writeResponse(to, r)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
return json.Marshal(out)
|
||||
}
|
||||
func readResponse(k Kind, m object) (response, error) {
|
||||
r := response{id: str(m["id"]), model: str(m["model"]), usage: usage(k, obj(m["usage"]))}
|
||||
if m["error"] != nil {
|
||||
return r, fmt.Errorf("upstream error: %v", m["error"])
|
||||
}
|
||||
var e error
|
||||
switch k {
|
||||
case Chat:
|
||||
if e = fields(m, "id", "object", "created", "model", "choices", "usage", "system_fingerprint", "service_tier"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
cs := arr(m["choices"])
|
||||
if len(cs) != 1 {
|
||||
return r, fmt.Errorf("expected one choice")
|
||||
}
|
||||
c := obj(cs[0])
|
||||
if e = fields(c, "index", "message", "finish_reason", "logprobs"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if c["logprobs"] != nil {
|
||||
return r, fmt.Errorf("logprobs unsupported")
|
||||
}
|
||||
message, e := chatResponseMessage(obj(c["message"]))
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
mm, e := readMessage(k, message)
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.parts = mm.parts
|
||||
r.finish, e = finishReason(k, str(c["finish_reason"]))
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
case Anthropic:
|
||||
if e = fields(m, "id", "type", "role", "model", "content", "stop_reason", "stop_sequence", "usage"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if m["stop_sequence"] != nil {
|
||||
return r, fmt.Errorf("stop sequence cannot be preserved")
|
||||
}
|
||||
content, e := anthropicOutput(m["content"])
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.parts, e = readContent(k, content)
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.finish, e = finishReason(k, str(m["stop_reason"]))
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
case Responses:
|
||||
// Responses echoes request configuration; only output items carry generated content.
|
||||
for _, v := range arr(m["output"]) {
|
||||
x := obj(v)
|
||||
switch str(x["type"]) {
|
||||
case "reasoning":
|
||||
if e = validateReasoning(x); e != nil {
|
||||
return r, e
|
||||
}
|
||||
case "message":
|
||||
x = responseMessage(x)
|
||||
mm, e := readMessage(k, x)
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.parts = append(r.parts, mm.parts...)
|
||||
case "function_call":
|
||||
if e = fields(x, "type", "id", "call_id", "name", "arguments", "status"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
p, e := callPart(x, "call_id", "arguments")
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.parts = append(r.parts, p)
|
||||
default:
|
||||
return r, fmt.Errorf("unsupported output item %q", str(x["type"]))
|
||||
}
|
||||
}
|
||||
if str(m["status"]) == "incomplete" {
|
||||
if str(obj(m["incomplete_details"])["reason"]) != "max_output_tokens" {
|
||||
return r, fmt.Errorf("unsupported incomplete reason")
|
||||
}
|
||||
r.finish = "length"
|
||||
} else {
|
||||
r.finish, e = finishReason(k, str(m["status"]))
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
for _, p := range r.parts {
|
||||
if p.typ == "call" {
|
||||
r.finish = "tool"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
func writeResponse(k Kind, r response) (object, error) {
|
||||
if r.id == "" {
|
||||
r.id = newID()
|
||||
}
|
||||
out := object{"id": r.id, "model": r.model}
|
||||
cs := []any{}
|
||||
calls := []any{}
|
||||
text := ""
|
||||
for i, p := range r.parts {
|
||||
if p.typ != "text" && p.typ != "call" {
|
||||
return nil, fmt.Errorf("unsupported response content")
|
||||
}
|
||||
b, e := contentBlock(k, p, "assistant")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
switch k {
|
||||
case Chat:
|
||||
if p.typ == "call" {
|
||||
calls = append(calls, b)
|
||||
} else {
|
||||
text += p.text
|
||||
}
|
||||
case Anthropic:
|
||||
cs = append(cs, b)
|
||||
case Responses:
|
||||
if p.typ == "call" {
|
||||
b["id"] = fmt.Sprintf("fc_%s_%d", r.id, i)
|
||||
b["status"] = "completed"
|
||||
cs = append(cs, b)
|
||||
} else {
|
||||
b["annotations"] = []any{}
|
||||
cs = append(cs, object{"type": "message", "id": fmt.Sprintf("msg_%s_%d", r.id, i), "role": "assistant", "status": "completed", "content": []any{b}})
|
||||
}
|
||||
}
|
||||
}
|
||||
switch k {
|
||||
case Chat:
|
||||
m := object{"role": "assistant", "content": text}
|
||||
if len(calls) > 0 {
|
||||
m["tool_calls"] = calls
|
||||
if text == "" {
|
||||
m["content"] = nil
|
||||
}
|
||||
}
|
||||
out["object"] = "chat.completion"
|
||||
out["created"] = 0
|
||||
out["choices"] = []any{object{"index": 0, "message": m, "finish_reason": wireFinish(k, r.finish)}}
|
||||
case Anthropic:
|
||||
out["type"] = "message"
|
||||
out["role"] = "assistant"
|
||||
out["content"] = cs
|
||||
out["stop_reason"] = wireFinish(k, r.finish)
|
||||
out["stop_sequence"] = nil
|
||||
case Responses:
|
||||
out["object"] = "response"
|
||||
out["created_at"] = time.Now().Unix()
|
||||
out["status"] = wireFinish(k, r.finish)
|
||||
out["output"] = cs
|
||||
out["error"] = nil
|
||||
out["incomplete_details"] = nil
|
||||
if r.finish == "length" {
|
||||
out["incomplete_details"] = object{"reason": "max_output_tokens"}
|
||||
}
|
||||
}
|
||||
if r.usage.Known || k == Anthropic {
|
||||
out["usage"] = usageJSON(k, r.usage)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,669 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
type Kind string
|
||||
|
||||
const (
|
||||
Chat Kind = "openai-chat"
|
||||
Responses Kind = "openai-responses"
|
||||
Anthropic Kind = "anthropic"
|
||||
)
|
||||
|
||||
type Usage struct {
|
||||
Input int64 `json:"input_tokens"`
|
||||
Output int64 `json:"output_tokens"`
|
||||
CacheRead int64 `json:"cache_read_tokens"`
|
||||
CacheWrite int64 `json:"cache_write_tokens"`
|
||||
Known bool `json:"known"`
|
||||
Raw json.RawMessage `json:"raw,omitempty"`
|
||||
}
|
||||
type object = map[string]any
|
||||
|
||||
// maxJSONNodes bounds how many JSON tokens one body may contain before it is
|
||||
// materialized into maps. Without it a 16 MiB body of micro-values expands
|
||||
// into hundreds of MB of Go objects per request.
|
||||
const maxJSONNodes = 1 << 18
|
||||
|
||||
func valid(k Kind) bool { return k == Chat || k == Responses || k == Anthropic }
|
||||
|
||||
// scanNodeCount walks b with the streaming decoder and fails once more than
|
||||
// limit tokens have been seen. Malformed JSON is not an error here; the
|
||||
// materializing pass reports its usual message.
|
||||
func scanNodeCount(b []byte, limit int) error {
|
||||
d := json.NewDecoder(bytes.NewReader(b))
|
||||
n := 0
|
||||
for {
|
||||
_, err := d.Token()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if n++; n > limit {
|
||||
return fmt.Errorf("JSON exceeds %d nodes", limit)
|
||||
}
|
||||
}
|
||||
}
|
||||
func decode(b []byte) (object, error) {
|
||||
if err := scanNodeCount(b, maxJSONNodes); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
d := json.NewDecoder(bytes.NewReader(b))
|
||||
d.UseNumber()
|
||||
var v object
|
||||
if err := d.Decode(&v); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if v == nil {
|
||||
return nil, fmt.Errorf("expected JSON object")
|
||||
}
|
||||
var extra any
|
||||
if err := d.Decode(&extra); err != io.EOF {
|
||||
return nil, fmt.Errorf("trailing JSON data")
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
func obj(v any) object { m, _ := v.(map[string]any); return m }
|
||||
func str(v any) string { s, _ := v.(string); return s }
|
||||
func arr(v any) []any { a, _ := v.([]any); return a }
|
||||
func num(v any) int64 {
|
||||
switch n := v.(type) {
|
||||
case json.Number:
|
||||
i, _ := n.Int64()
|
||||
return i
|
||||
case int64:
|
||||
return n
|
||||
case int:
|
||||
return int64(n)
|
||||
case float64:
|
||||
return int64(n)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
func fields(m object, names ...string) error {
|
||||
allowed := map[string]bool{}
|
||||
for _, n := range names {
|
||||
allowed[n] = true
|
||||
}
|
||||
for k := range m {
|
||||
if !allowed[k] {
|
||||
return fmt.Errorf("unsupported field %q", k)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func required(m object, k string) (string, error) {
|
||||
s, ok := m[k].(string)
|
||||
if !ok || s == "" {
|
||||
return "", fmt.Errorf("missing or invalid %s", k)
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
func ExtractUsage(kind Kind, body []byte) Usage {
|
||||
m, e := decode(body)
|
||||
if e != nil {
|
||||
return Usage{}
|
||||
}
|
||||
if r := obj(m["response"]); r != nil {
|
||||
m = r
|
||||
}
|
||||
if r := obj(m["message"]); r != nil {
|
||||
m = r
|
||||
}
|
||||
return usage(kind, obj(m["usage"]))
|
||||
}
|
||||
func usage(k Kind, m object) Usage {
|
||||
if m == nil {
|
||||
return Usage{}
|
||||
}
|
||||
raw, _ := json.Marshal(m)
|
||||
u := Usage{Raw: raw}
|
||||
switch k {
|
||||
case Chat:
|
||||
_, a := m["prompt_tokens"]
|
||||
_, b := m["completion_tokens"]
|
||||
u.Known = a || b
|
||||
u.Input = num(m["prompt_tokens"])
|
||||
u.Output = num(m["completion_tokens"])
|
||||
u.CacheRead = num(obj(m["prompt_tokens_details"])["cached_tokens"])
|
||||
case Responses:
|
||||
_, a := m["input_tokens"]
|
||||
_, b := m["output_tokens"]
|
||||
u.Known = a || b
|
||||
u.Input = num(m["input_tokens"])
|
||||
u.Output = num(m["output_tokens"])
|
||||
u.CacheRead = num(obj(m["input_tokens_details"])["cached_tokens"])
|
||||
case Anthropic:
|
||||
_, a := m["input_tokens"]
|
||||
_, b := m["output_tokens"]
|
||||
u.Known = a || b
|
||||
u.CacheRead = num(m["cache_read_input_tokens"])
|
||||
u.CacheWrite = num(m["cache_creation_input_tokens"])
|
||||
u.Input = num(m["input_tokens"]) + u.CacheRead + u.CacheWrite
|
||||
u.Output = num(m["output_tokens"])
|
||||
}
|
||||
return u
|
||||
}
|
||||
func usageJSON(k Kind, u Usage) object {
|
||||
switch k {
|
||||
case Chat:
|
||||
return object{"prompt_tokens": u.Input, "completion_tokens": u.Output, "total_tokens": u.Input + u.Output, "prompt_tokens_details": object{"cached_tokens": u.CacheRead}}
|
||||
case Responses:
|
||||
return object{"input_tokens": u.Input, "output_tokens": u.Output, "total_tokens": u.Input + u.Output, "input_tokens_details": object{"cached_tokens": u.CacheRead}}
|
||||
default:
|
||||
return object{"input_tokens": u.Input - u.CacheRead - u.CacheWrite, "output_tokens": u.Output, "cache_read_input_tokens": u.CacheRead, "cache_creation_input_tokens": u.CacheWrite}
|
||||
}
|
||||
}
|
||||
|
||||
type part struct{ typ, text, id, name, args, url string }
|
||||
type message struct {
|
||||
role string
|
||||
parts []part
|
||||
}
|
||||
type request struct {
|
||||
messages []message
|
||||
tools []object
|
||||
options object
|
||||
choice any
|
||||
}
|
||||
|
||||
func textPart(s string) part { return part{typ: "text", text: s} }
|
||||
|
||||
func ConvertRequest(from, to Kind, body []byte, model string) ([]byte, error) {
|
||||
if !valid(from) || !valid(to) {
|
||||
return nil, fmt.Errorf("unknown protocol")
|
||||
}
|
||||
m, err := decode(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if from == to {
|
||||
if model != "" {
|
||||
m["model"] = model
|
||||
}
|
||||
return json.Marshal(m)
|
||||
}
|
||||
r, err := readRequest(from, m)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, err := writeRequest(to, r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if model == "" {
|
||||
model = str(m["model"])
|
||||
}
|
||||
out["model"] = model
|
||||
return json.Marshal(out)
|
||||
}
|
||||
func readRequest(k Kind, m object) (request, error) {
|
||||
r := request{options: object{}}
|
||||
var err error
|
||||
common := []string{"model", "stream", "temperature", "top_p", "tools", "tool_choice"}
|
||||
switch k {
|
||||
case Chat:
|
||||
err = fields(m, append(common, "messages", "max_tokens", "max_completion_tokens", "stop", "stream_options", "parallel_tool_calls", "n")...)
|
||||
case Responses:
|
||||
err = fields(m, append(common, "input", "instructions", "max_output_tokens", "parallel_tool_calls", "store")...)
|
||||
case Anthropic:
|
||||
err = fields(m, append(common, "messages", "system", "max_tokens", "stop_sequences")...)
|
||||
}
|
||||
if err != nil {
|
||||
return r, err
|
||||
}
|
||||
if n, ok := m["n"]; ok && num(n) != 1 {
|
||||
return r, fmt.Errorf("multiple choices unsupported")
|
||||
}
|
||||
if v, ok := m["store"]; ok && v != false {
|
||||
return r, fmt.Errorf("stored Responses unsupported")
|
||||
}
|
||||
if v, ok := m["stream_options"]; ok {
|
||||
o := obj(v)
|
||||
if o == nil {
|
||||
return r, fmt.Errorf("invalid stream_options")
|
||||
}
|
||||
if e := fields(o, "include_usage"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"stream", "temperature", "top_p", "parallel_tool_calls"} {
|
||||
if v, ok := m[key]; ok {
|
||||
r.options[key] = v
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens"} {
|
||||
if v, ok := m[key]; ok {
|
||||
if _, exists := r.options["max"]; exists {
|
||||
return r, fmt.Errorf("conflicting token limits")
|
||||
}
|
||||
r.options["max"] = v
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"stop", "stop_sequences"} {
|
||||
if v, ok := m[key]; ok {
|
||||
if s, ok := v.(string); ok {
|
||||
v = []any{s}
|
||||
}
|
||||
if _, ok := v.([]any); !ok {
|
||||
return r, fmt.Errorf("invalid stop sequences")
|
||||
}
|
||||
r.options["stop"] = v
|
||||
}
|
||||
}
|
||||
if k == Responses {
|
||||
if s, ok := m["instructions"]; ok {
|
||||
v, ok := s.(string)
|
||||
if !ok {
|
||||
return r, fmt.Errorf("invalid instructions")
|
||||
}
|
||||
r.messages = append(r.messages, message{"system", []part{textPart(v)}})
|
||||
}
|
||||
if s, ok := m["input"].(string); ok {
|
||||
r.messages = append(r.messages, message{"user", []part{textPart(s)}})
|
||||
} else {
|
||||
items, ok := m["input"].([]any)
|
||||
if !ok {
|
||||
return r, fmt.Errorf("invalid input")
|
||||
}
|
||||
for _, v := range items {
|
||||
x := obj(v)
|
||||
typ := str(x["type"])
|
||||
switch typ {
|
||||
case "function_call":
|
||||
if e := fields(x, "type", "id", "call_id", "name", "arguments", "status"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
p, e := callPart(x, "call_id", "arguments")
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.messages = append(r.messages, message{"assistant", []part{p}})
|
||||
case "function_call_output":
|
||||
if e := fields(x, "type", "id", "call_id", "output", "status"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
id, e := required(x, "call_id")
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
s, ok := x["output"].(string)
|
||||
if !ok {
|
||||
return r, fmt.Errorf("non-text function output unsupported")
|
||||
}
|
||||
r.messages = append(r.messages, message{"tool", []part{{typ: "result", id: id, text: s}}})
|
||||
case "", "message":
|
||||
mm, e := readMessage(k, x)
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.messages = append(r.messages, mm)
|
||||
default:
|
||||
return r, fmt.Errorf("unsupported input item %q", typ)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if k == Anthropic {
|
||||
if s, ok := m["system"]; ok {
|
||||
ps, e := readContent(k, s)
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
for _, p := range ps {
|
||||
if p.typ != "text" {
|
||||
return r, fmt.Errorf("non-text system")
|
||||
}
|
||||
}
|
||||
r.messages = append(r.messages, message{"system", ps})
|
||||
}
|
||||
}
|
||||
ms, ok := m["messages"].([]any)
|
||||
if !ok {
|
||||
return r, fmt.Errorf("invalid messages")
|
||||
}
|
||||
for _, v := range ms {
|
||||
mm, e := readMessage(k, obj(v))
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.messages = append(r.messages, mm)
|
||||
}
|
||||
}
|
||||
if v, ok := m["tools"]; ok {
|
||||
ts, ok := v.([]any)
|
||||
if !ok {
|
||||
return r, fmt.Errorf("invalid tools")
|
||||
}
|
||||
for _, v := range ts {
|
||||
t := obj(v)
|
||||
var f object
|
||||
switch k {
|
||||
case Chat:
|
||||
if e := fields(t, "type", "function"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if str(t["type"]) != "function" {
|
||||
return r, fmt.Errorf("unsupported tool")
|
||||
}
|
||||
f = obj(t["function"])
|
||||
case Responses:
|
||||
if str(t["type"]) != "function" {
|
||||
return r, fmt.Errorf("unsupported tool")
|
||||
}
|
||||
f = object{}
|
||||
for key, val := range t {
|
||||
if key != "type" {
|
||||
f[key] = val
|
||||
}
|
||||
}
|
||||
case Anthropic:
|
||||
f = object{}
|
||||
for key, val := range t {
|
||||
if key == "input_schema" {
|
||||
key = "parameters"
|
||||
}
|
||||
f[key] = val
|
||||
}
|
||||
}
|
||||
if e := fields(f, "name", "description", "parameters", "strict"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if _, e := required(f, "name"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if obj(f["parameters"]) == nil {
|
||||
return r, fmt.Errorf("missing tool schema")
|
||||
}
|
||||
r.tools = append(r.tools, f)
|
||||
}
|
||||
}
|
||||
if v, ok := m["tool_choice"]; ok {
|
||||
if s, ok := v.(string); ok {
|
||||
if s != "auto" && s != "none" && s != "required" {
|
||||
return r, fmt.Errorf("unsupported tool_choice")
|
||||
}
|
||||
r.choice = s
|
||||
} else {
|
||||
x := obj(v)
|
||||
if k == Chat {
|
||||
if e := fields(x, "type", "function"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if str(x["type"]) != "function" {
|
||||
return r, fmt.Errorf("unsupported tool_choice")
|
||||
}
|
||||
f := obj(x["function"])
|
||||
if e := fields(f, "name"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
s, e := required(f, "name")
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.choice = object{"name": s}
|
||||
} else if k == Responses {
|
||||
if e := fields(x, "type", "name"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
if str(x["type"]) != "function" {
|
||||
return r, fmt.Errorf("unsupported tool_choice")
|
||||
}
|
||||
s, e := required(x, "name")
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.choice = object{"name": s}
|
||||
} else {
|
||||
if e := fields(x, "type", "name", "disable_parallel_tool_use"); e != nil {
|
||||
return r, e
|
||||
}
|
||||
switch str(x["type"]) {
|
||||
case "auto", "none":
|
||||
r.choice = x["type"]
|
||||
case "any":
|
||||
r.choice = "required"
|
||||
case "tool":
|
||||
s, e := required(x, "name")
|
||||
if e != nil {
|
||||
return r, e
|
||||
}
|
||||
r.choice = object{"name": s}
|
||||
default:
|
||||
return r, fmt.Errorf("unsupported tool_choice")
|
||||
}
|
||||
if b, ok := x["disable_parallel_tool_use"].(bool); ok {
|
||||
r.options["parallel_tool_calls"] = !b
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
func callPart(x object, idKey, argsKey string) (part, error) {
|
||||
id, e := required(x, idKey)
|
||||
if e != nil {
|
||||
return part{}, e
|
||||
}
|
||||
name, e := required(x, "name")
|
||||
if e != nil {
|
||||
return part{}, e
|
||||
}
|
||||
s, ok := x[argsKey].(string)
|
||||
if !ok {
|
||||
return part{}, fmt.Errorf("invalid tool arguments")
|
||||
}
|
||||
if _, e := decode([]byte(s)); e != nil {
|
||||
return part{}, fmt.Errorf("tool arguments must be JSON object: %w", e)
|
||||
}
|
||||
return part{typ: "call", id: id, name: name, args: s}, nil
|
||||
}
|
||||
func readMessage(k Kind, x object) (message, error) {
|
||||
m := message{role: str(x["role"])}
|
||||
if x == nil {
|
||||
return m, fmt.Errorf("invalid message")
|
||||
}
|
||||
if e := fields(x, "role", "content", "type", "id", "status", "tool_calls", "tool_call_id", "refusal"); e != nil {
|
||||
return m, e
|
||||
}
|
||||
if x["refusal"] != nil {
|
||||
return m, fmt.Errorf("refusal content cannot be preserved")
|
||||
}
|
||||
switch m.role {
|
||||
case "user", "assistant", "system", "developer", "tool":
|
||||
default:
|
||||
return m, fmt.Errorf("unsupported role %q", m.role)
|
||||
}
|
||||
if m.role == "tool" {
|
||||
if k != Chat {
|
||||
return m, fmt.Errorf("invalid tool role")
|
||||
}
|
||||
id, e := required(x, "tool_call_id")
|
||||
if e != nil {
|
||||
return m, e
|
||||
}
|
||||
s, ok := x["content"].(string)
|
||||
if !ok {
|
||||
return m, fmt.Errorf("non-text tool result unsupported")
|
||||
}
|
||||
m.parts = []part{{typ: "result", id: id, text: s}}
|
||||
return m, nil
|
||||
}
|
||||
if c := x["content"]; c != nil {
|
||||
ps, e := readContent(k, c)
|
||||
if e != nil {
|
||||
return m, e
|
||||
}
|
||||
m.parts = ps
|
||||
}
|
||||
if tc, ok := x["tool_calls"]; ok {
|
||||
if k != Chat || m.role != "assistant" {
|
||||
return m, fmt.Errorf("invalid tool_calls")
|
||||
}
|
||||
ts, ok := tc.([]any)
|
||||
if !ok {
|
||||
return m, fmt.Errorf("invalid tool_calls")
|
||||
}
|
||||
for _, v := range ts {
|
||||
t := obj(v)
|
||||
if e := fields(t, "id", "type", "function"); e != nil {
|
||||
return m, e
|
||||
}
|
||||
if str(t["type"]) != "function" {
|
||||
return m, fmt.Errorf("unsupported tool call")
|
||||
}
|
||||
f := obj(t["function"])
|
||||
if e := fields(f, "name", "arguments"); e != nil {
|
||||
return m, e
|
||||
}
|
||||
p, e := callPart(object{"id": t["id"], "name": f["name"], "arguments": f["arguments"]}, "id", "arguments")
|
||||
if e != nil {
|
||||
return m, e
|
||||
}
|
||||
m.parts = append(m.parts, p)
|
||||
}
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
func readContent(k Kind, v any) ([]part, error) {
|
||||
if s, ok := v.(string); ok {
|
||||
return []part{textPart(s)}, nil
|
||||
}
|
||||
a, ok := v.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid content")
|
||||
}
|
||||
var ps []part
|
||||
for _, v := range a {
|
||||
x := obj(v)
|
||||
switch str(x["type"]) {
|
||||
case "text", "input_text", "output_text":
|
||||
if e := fields(x, "type", "text", "annotations"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if an, ok := x["annotations"]; ok && len(arr(an)) != 0 {
|
||||
return nil, fmt.Errorf("annotations unsupported")
|
||||
}
|
||||
s, ok := x["text"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid text")
|
||||
}
|
||||
ps = append(ps, textPart(s))
|
||||
case "image_url":
|
||||
if k != Chat {
|
||||
return nil, fmt.Errorf("invalid image block")
|
||||
}
|
||||
if e := fields(x, "type", "image_url"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
im := obj(x["image_url"])
|
||||
if e := fields(im, "url", "detail"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if d := str(im["detail"]); d != "" && d != "auto" {
|
||||
return nil, fmt.Errorf("image detail cannot be preserved")
|
||||
}
|
||||
s, e := required(im, "url")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
ps = append(ps, part{typ: "image", url: s})
|
||||
case "input_image":
|
||||
if k != Responses {
|
||||
return nil, fmt.Errorf("invalid image block")
|
||||
}
|
||||
if e := fields(x, "type", "image_url", "detail"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if d := str(x["detail"]); d != "" && d != "auto" {
|
||||
return nil, fmt.Errorf("image detail cannot be preserved")
|
||||
}
|
||||
s, e := required(x, "image_url")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
ps = append(ps, part{typ: "image", url: s})
|
||||
case "image":
|
||||
if k != Anthropic {
|
||||
return nil, fmt.Errorf("invalid image block")
|
||||
}
|
||||
if e := fields(x, "type", "source"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
src := obj(x["source"])
|
||||
var url string
|
||||
switch str(src["type"]) {
|
||||
case "url":
|
||||
if e := fields(src, "type", "url"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
url = str(src["url"])
|
||||
case "base64":
|
||||
if e := fields(src, "type", "media_type", "data"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if str(src["media_type"]) == "" || str(src["data"]) == "" {
|
||||
return nil, fmt.Errorf("invalid image source")
|
||||
}
|
||||
url = "data:" + str(src["media_type"]) + ";base64," + str(src["data"])
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported image source")
|
||||
}
|
||||
if url == "" {
|
||||
return nil, fmt.Errorf("empty image URL")
|
||||
}
|
||||
ps = append(ps, part{typ: "image", url: url})
|
||||
case "tool_use":
|
||||
if k != Anthropic {
|
||||
return nil, fmt.Errorf("invalid tool block")
|
||||
}
|
||||
if e := fields(x, "type", "id", "name", "input"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
b, e := json.Marshal(x["input"])
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
p, e := callPart(object{"id": x["id"], "name": x["name"], "arguments": string(b)}, "id", "arguments")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
ps = append(ps, p)
|
||||
case "tool_result":
|
||||
if k != Anthropic {
|
||||
return nil, fmt.Errorf("invalid result block")
|
||||
}
|
||||
if e := fields(x, "type", "tool_use_id", "content", "is_error"); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if x["is_error"] == true {
|
||||
return nil, fmt.Errorf("tool error flag cannot be preserved")
|
||||
}
|
||||
id, e := required(x, "tool_use_id")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
s, ok := x["content"].(string)
|
||||
if !ok {
|
||||
pp, e := readContent(k, x["content"])
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
for _, p := range pp {
|
||||
if p.typ != "text" {
|
||||
return nil, fmt.Errorf("non-text tool result unsupported")
|
||||
}
|
||||
s += p.text
|
||||
}
|
||||
}
|
||||
ps = append(ps, part{typ: "result", id: id, text: s})
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported content type %q", str(x["type"]))
|
||||
}
|
||||
}
|
||||
return ps, nil
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRequestMatrix(t *testing.T) {
|
||||
chat := []byte(`{"model":"old","max_tokens":100,"stream":true,"messages":[{"role":"system","content":"help"},{"role":"user","content":[{"type":"text","text":"look"},{"type":"image_url","image_url":{"url":"data:image/png;base64,YQ=="}}]},{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"read","arguments":"{\"path\":\"a\"}"}}]},{"role":"tool","tool_call_id":"c1","content":"ok"},{"role":"user","content":"continue"}],"tools":[{"type":"function","function":{"name":"read","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}}],"tool_choice":"auto"}`)
|
||||
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||
src := chat
|
||||
var e error
|
||||
if from != Chat {
|
||||
src, e = ConvertRequest(Chat, from, chat, "new")
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
}
|
||||
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||
t.Run(string(from)+"/"+string(to), func(t *testing.T) {
|
||||
b, e := ConvertRequest(from, to, src, "target")
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
m, e := decode(b)
|
||||
if e != nil || m["model"] != "target" {
|
||||
t.Fatalf("%s %v", b, e)
|
||||
}
|
||||
r, e := readRequest(to, m)
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
var calls, results, images int
|
||||
for _, m := range r.messages {
|
||||
for _, p := range m.parts {
|
||||
switch p.typ {
|
||||
case "call":
|
||||
calls++
|
||||
if p.id != "c1" {
|
||||
t.Fatal(p)
|
||||
}
|
||||
case "result":
|
||||
results++
|
||||
case "image":
|
||||
images++
|
||||
}
|
||||
}
|
||||
}
|
||||
if calls != 1 || results != 1 || images != 1 {
|
||||
t.Fatalf("lost content: %s", b)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
func TestStrictRejection(t *testing.T) {
|
||||
cases := []struct {
|
||||
k Kind
|
||||
b string
|
||||
to Kind
|
||||
}{
|
||||
{Chat, `{"messages":[],"response_format":{"type":"json_object"}}`, Responses},
|
||||
{Chat, `{"messages":[]}`, Anthropic},
|
||||
{Responses, `{"input":[],"previous_response_id":"r"}`, Chat},
|
||||
{Anthropic, `{"max_tokens":10,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x"}]}]}`, Chat},
|
||||
{Chat, `{"messages":[],"n":2}`, Responses},
|
||||
{Chat, `{"messages":[],"stop":["x"]}`, Responses},
|
||||
{Chat, `{"max_tokens":10,"messages":[],"tools":[{"type":"function","function":{"name":"x","parameters":{},"strict":true}}]}`, Anthropic},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if b, e := ConvertRequest(c.k, c.to, []byte(c.b), "m"); e == nil {
|
||||
t.Fatalf("accepted %s => %s", c.b, b)
|
||||
}
|
||||
}
|
||||
}
|
||||
func TestResponseMatrix(t *testing.T) {
|
||||
chat := []byte(`{"id":"r","model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"working","tool_calls":[{"id":"c","type":"function","function":{"name":"read","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":4,"prompt_tokens_details":{"cached_tokens":3}}}`)
|
||||
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||
src, e := ConvertResponse(Chat, from, chat)
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||
b, e := ConvertResponse(from, to, src)
|
||||
if e != nil {
|
||||
t.Fatal(from, to, e)
|
||||
}
|
||||
m, _ := decode(b)
|
||||
r, e := readResponse(to, m)
|
||||
if e != nil || len(r.parts) != 2 || r.finish != "tool" {
|
||||
t.Fatalf("%s %v", b, e)
|
||||
}
|
||||
u := ExtractUsage(to, b)
|
||||
if u.Input != 12 || u.Output != 4 || u.CacheRead != 3 || !u.Known {
|
||||
t.Fatal(u)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type flushBuffer struct {
|
||||
bytes.Buffer
|
||||
flushes int
|
||||
}
|
||||
|
||||
func (b *flushBuffer) Flush() { b.flushes++ }
|
||||
func chatStream() string {
|
||||
return strings.Join([]string{
|
||||
`data: {"id":"r","model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}`,
|
||||
`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"read","arguments":"{\"x\":"}},{"index":1,"id":"c2","type":"function","function":{"name":"write","arguments":"{"}}]}}]}`,
|
||||
`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"}"}},{"index":0,"function":{"arguments":"1}"}}]}}]}`,
|
||||
`data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`,
|
||||
`data: {"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":7}}`,
|
||||
`data: [DONE]`,
|
||||
}, "\n\n") + "\n\n"
|
||||
}
|
||||
func TestStreamMatrix(t *testing.T) {
|
||||
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||
src := chatStream()
|
||||
if from != Chat {
|
||||
var b bytes.Buffer
|
||||
if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, from, "m"); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
src = b.String()
|
||||
}
|
||||
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||
t.Run(string(from)+"/"+string(to), func(t *testing.T) {
|
||||
var b flushBuffer
|
||||
u, e := Stream(context.Background(), &b, strings.NewReader(src), from, to, "m")
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if u.Input != 10 || u.Output != 7 || !u.Known {
|
||||
t.Fatal(u)
|
||||
}
|
||||
if b.flushes < 3 {
|
||||
t.Fatal("not flushed")
|
||||
}
|
||||
if from == to && b.String() != src {
|
||||
t.Fatal("passthrough changed bytes")
|
||||
}
|
||||
if to == Responses {
|
||||
for _, event := range []string{"response.output_item.added", "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.completed"} {
|
||||
if !strings.Contains(b.String(), event) {
|
||||
t.Fatal("missing", event)
|
||||
}
|
||||
}
|
||||
}
|
||||
if to == Anthropic && !strings.Contains(b.String(), "message_stop") {
|
||||
t.Fatal("missing stop")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
func TestUsageCumulative(t *testing.T) {
|
||||
src := "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":5,\"output_tokens\":1,\"cache_read_input_tokens\":8,\"cache_creation_input_tokens\":2}}}\n\n" + "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":9}}\n\n"
|
||||
var b bytes.Buffer
|
||||
u, e := Stream(context.Background(), &b, strings.NewReader(src), Anthropic, Anthropic, "")
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if u.Input != 15 || u.Output != 9 || u.CacheRead != 8 || u.CacheWrite != 2 {
|
||||
t.Fatal(u)
|
||||
}
|
||||
raw, e := decode(u.Raw)
|
||||
if e != nil || num(raw["input_tokens"]) != 5 || num(raw["output_tokens"]) != 9 || num(raw["cache_read_input_tokens"]) != 8 {
|
||||
t.Fatal(string(u.Raw), e)
|
||||
}
|
||||
}
|
||||
func TestRealtimeBeforeEOF(t *testing.T) {
|
||||
r, w := io.Pipe()
|
||||
defer r.Close()
|
||||
defer w.Close()
|
||||
done := make(chan error, 1)
|
||||
out := &signalWriter{signal: make(chan struct{}, 1)}
|
||||
go func() { _, e := Stream(context.Background(), out, r, Chat, Responses, "m"); done <- e }()
|
||||
first, rest, _ := strings.Cut(chatStream(), "\n\n")
|
||||
if _, e := io.WriteString(w, first+"\n\n"); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
select {
|
||||
case <-out.signal:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first delta buffered until EOF")
|
||||
}
|
||||
if _, e := io.WriteString(w, rest); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
select {
|
||||
case e := <-done:
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("terminal marker waits for EOF")
|
||||
}
|
||||
}
|
||||
|
||||
type signalWriter struct{ signal chan struct{} }
|
||||
|
||||
func (w *signalWriter) Write(p []byte) (int, error) {
|
||||
select {
|
||||
case w.signal <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
func TestUnknownUsageAndUniqueID(t *testing.T) {
|
||||
src := "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"a\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"
|
||||
var previous string
|
||||
for i := 0; i < 2; i++ {
|
||||
var b bytes.Buffer
|
||||
u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Anthropic, "m")
|
||||
if e != nil || u.Known {
|
||||
t.Fatal(u, e)
|
||||
}
|
||||
line := strings.Split(b.String(), "\n")[1]
|
||||
m, e := decode([]byte(strings.TrimPrefix(line, "data: ")))
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
msg := obj(m["message"])
|
||||
if obj(msg["usage"]) == nil {
|
||||
t.Fatal("missing initial usage")
|
||||
}
|
||||
id := str(msg["id"])
|
||||
if id == "" || id == previous {
|
||||
t.Fatal("nonunique id")
|
||||
}
|
||||
previous = id
|
||||
}
|
||||
}
|
||||
func TestStreamErrors(t *testing.T) {
|
||||
for _, src := range []string{"data: nope\n\n", "data: {\"error\":{\"message\":\"bad\"}}\n\n", strings.Replace(chatStream(), "data: [DONE]\n\n", "", 1), strings.Replace(chatStream(), `"arguments":"1}"`, `"arguments":"oops}"`, 1)} {
|
||||
var b bytes.Buffer
|
||||
if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Responses, "m"); e == nil {
|
||||
t.Fatal("accepted broken stream", src)
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if _, e := Stream(ctx, io.Discard, strings.NewReader(chatStream()), Chat, Chat, ""); !errors.Is(e, context.Canceled) {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if _, e := Stream(context.Background(), badWriter{}, strings.NewReader(chatStream()), Chat, Anthropic, ""); !errors.Is(e, io.ErrClosedPipe) {
|
||||
t.Fatal(e)
|
||||
}
|
||||
}
|
||||
|
||||
type badWriter struct{}
|
||||
|
||||
func (badWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }
|
||||
func TestSSEMultilineAndCRLF(t *testing.T) {
|
||||
src := ": comment\r\n\r\ndata: {\r\ndata: \"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3}}\r\n\r\ndata: [DONE]\r\n\r\n"
|
||||
var b flushBuffer
|
||||
u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Chat, "")
|
||||
if e != nil || u.Input != 2 || u.Output != 3 || b.String() != src {
|
||||
t.Fatal(u, e, b.String())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Recognized Anthropic thinking blocks are omitted only from converted output.
|
||||
func anthropicThinking(x object) (bool, error) {
|
||||
var names []string
|
||||
switch str(x["type"]) {
|
||||
case "thinking":
|
||||
names = []string{"thinking", "signature"}
|
||||
case "redacted_thinking":
|
||||
names = []string{"data"}
|
||||
default:
|
||||
return false, nil
|
||||
}
|
||||
if err := fields(x, append([]string{"type"}, names...)...); err != nil {
|
||||
return true, err
|
||||
}
|
||||
for _, name := range names {
|
||||
if _, ok := x[name].(string); !ok {
|
||||
return true, fmt.Errorf("invalid thinking %s", name)
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func anthropicOutput(v any) (any, error) {
|
||||
content, ok := v.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid Anthropic content")
|
||||
}
|
||||
clean := []any{}
|
||||
for _, block := range content {
|
||||
thinking, err := anthropicThinking(obj(block))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !thinking {
|
||||
clean = append(clean, block)
|
||||
}
|
||||
}
|
||||
return clean, nil
|
||||
}
|
||||
|
||||
type thinkingBlock struct {
|
||||
kind string
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (s *streamState) thinkingEvent(typ string, m object) (bool, error) {
|
||||
if typ == "message_delta" || typ == "message_stop" {
|
||||
for _, b := range s.thinking {
|
||||
if !b.closed {
|
||||
return true, fmt.Errorf("unfinished thinking block")
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
if typ != "content_block_start" && typ != "content_block_delta" && typ != "content_block_stop" {
|
||||
return false, nil
|
||||
}
|
||||
index, err := responseIndex(m, "index")
|
||||
if err != nil {
|
||||
return true, err
|
||||
}
|
||||
b := s.thinking[index]
|
||||
if typ == "content_block_start" {
|
||||
if b != nil || s.source[index] != nil {
|
||||
return true, fmt.Errorf("duplicate content index")
|
||||
}
|
||||
x := obj(m["content_block"])
|
||||
thinking, err := anthropicThinking(x)
|
||||
if err != nil || !thinking {
|
||||
return thinking, err
|
||||
}
|
||||
if !s.started || s.finished {
|
||||
return true, fmt.Errorf("thinking outside active message")
|
||||
}
|
||||
if s.thinking == nil {
|
||||
s.thinking = map[int]*thinkingBlock{}
|
||||
}
|
||||
s.thinking[index] = &thinkingBlock{kind: str(x["type"])}
|
||||
return true, nil
|
||||
}
|
||||
if b == nil {
|
||||
return false, nil
|
||||
}
|
||||
if b.closed || s.finished {
|
||||
return true, fmt.Errorf("event after thinking block closed")
|
||||
}
|
||||
if typ == "content_block_stop" {
|
||||
b.closed = true
|
||||
return true, nil
|
||||
}
|
||||
d := obj(m["delta"])
|
||||
field := ""
|
||||
switch str(d["type"]) {
|
||||
case "thinking_delta":
|
||||
field = "thinking"
|
||||
case "signature_delta":
|
||||
field = "signature"
|
||||
}
|
||||
if b.kind != "thinking" || field == "" {
|
||||
return true, fmt.Errorf("invalid thinking delta")
|
||||
}
|
||||
if err := fields(d, "type", field); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if _, ok := d[field].(string); !ok {
|
||||
return true, fmt.Errorf("invalid thinking delta value")
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// chatOutput omits output-only reasoning without exposing it as visible text.
|
||||
func chatOutput(x object) (object, error) {
|
||||
if v := x["reasoning_content"]; v != nil {
|
||||
if _, ok := v.(string); !ok {
|
||||
return nil, fmt.Errorf("reasoning_content must be a string or null")
|
||||
}
|
||||
}
|
||||
m := object{}
|
||||
for k, v := range x {
|
||||
if k != "reasoning_content" {
|
||||
m[k] = v
|
||||
}
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// Some Chat providers include streaming-style indices in complete tool calls.
|
||||
func chatResponseMessage(x object) (object, error) {
|
||||
m, err := chatOutput(x)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if calls, ok := x["tool_calls"].([]any); ok {
|
||||
clean := make([]any, len(calls))
|
||||
for i, v := range calls {
|
||||
call := obj(v)
|
||||
if _, exists := call["index"]; exists {
|
||||
index, err := responseIndex(call, "index")
|
||||
if err != nil || index != i {
|
||||
return nil, fmt.Errorf("invalid nonstream tool index")
|
||||
}
|
||||
}
|
||||
c := object{}
|
||||
for k, v := range call {
|
||||
if k != "index" {
|
||||
c[k] = v
|
||||
}
|
||||
}
|
||||
clean[i] = c
|
||||
}
|
||||
m["tool_calls"] = clean
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// Phase describes output metadata, not message content or request semantics.
|
||||
func responseMessage(x object) object {
|
||||
m := object{}
|
||||
for k, v := range x {
|
||||
if k != "phase" {
|
||||
m[k] = v
|
||||
}
|
||||
}
|
||||
if content, ok := x["content"].([]any); ok {
|
||||
clean := make([]any, len(content))
|
||||
for i, v := range content {
|
||||
p := obj(v)
|
||||
if logs, ok := p["logprobs"].([]any); ok && len(logs) == 0 && str(p["type"]) == "output_text" {
|
||||
q := object{}
|
||||
for k, v := range p {
|
||||
if k != "logprobs" {
|
||||
q[k] = v
|
||||
}
|
||||
}
|
||||
clean[i] = q
|
||||
} else {
|
||||
clean[i] = v
|
||||
}
|
||||
}
|
||||
m["content"] = clean
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func validateReasoning(x object) error {
|
||||
if err := fields(x, "type", "id", "status", "content", "summary", "encrypted_content"); err != nil {
|
||||
return err
|
||||
}
|
||||
for field, typ := range map[string]string{"content": "reasoning_text", "summary": "summary_text"} {
|
||||
if v, exists := x[field]; exists && v != nil {
|
||||
if _, ok := v.([]any); !ok {
|
||||
return fmt.Errorf("reasoning %s must be an array", field)
|
||||
}
|
||||
}
|
||||
for _, v := range arr(x[field]) {
|
||||
p := obj(v)
|
||||
if str(p["type"]) != typ {
|
||||
return fmt.Errorf("unsupported reasoning %s content", field)
|
||||
}
|
||||
if err := fields(p, "type", "text"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type reasoningItem struct {
|
||||
id string
|
||||
closed bool
|
||||
parts map[string]bool
|
||||
}
|
||||
|
||||
func responseIndex(m object, field string) (int, error) {
|
||||
n, ok := m[field].(json.Number)
|
||||
i, err := n.Int64()
|
||||
if !ok || err != nil || i < 0 || int64(int(i)) != i {
|
||||
return 0, fmt.Errorf("invalid %s", field)
|
||||
}
|
||||
return int(i), nil
|
||||
}
|
||||
|
||||
func (s *streamState) reasoningEvent(typ string, m object) (bool, error) {
|
||||
if typ == "response.completed" || typ == "response.incomplete" {
|
||||
seen := 0
|
||||
for i, v := range arr(obj(m["response"])["output"]) {
|
||||
x := obj(v)
|
||||
r := s.reasoning[i]
|
||||
if str(x["type"]) == "reasoning" {
|
||||
if r == nil || !r.closed || r.id != str(x["id"]) {
|
||||
return true, fmt.Errorf("final reasoning item disagrees with stream")
|
||||
}
|
||||
seen++
|
||||
} else if r != nil {
|
||||
return true, fmt.Errorf("final reasoning index changed type")
|
||||
}
|
||||
}
|
||||
if seen != len(s.reasoning) {
|
||||
return true, fmt.Errorf("final response missing reasoning item")
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
known := false
|
||||
switch typ {
|
||||
case "response.reasoning_text.delta", "response.reasoning_text.done",
|
||||
"response.reasoning_summary_text.delta", "response.reasoning_summary_text.done",
|
||||
"response.reasoning_summary_part.added", "response.reasoning_summary_part.done":
|
||||
known = true
|
||||
case "response.output_item.added", "response.output_item.done",
|
||||
"response.content_part.added", "response.content_part.done",
|
||||
"response.output_text.delta", "response.output_text.done",
|
||||
"response.function_call_arguments.delta", "response.function_call_arguments.done":
|
||||
default:
|
||||
return false, nil
|
||||
}
|
||||
key, err := responseIndex(m, "output_index")
|
||||
if err != nil {
|
||||
return true, err
|
||||
}
|
||||
r := s.reasoning[key]
|
||||
x := obj(m["item"])
|
||||
if typ == "response.output_item.added" {
|
||||
if r != nil || s.source[key] != nil {
|
||||
return true, fmt.Errorf("duplicate output index %d", key)
|
||||
}
|
||||
if str(x["type"]) != "reasoning" {
|
||||
return false, nil
|
||||
}
|
||||
if err := validateReasoning(x); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if str(x["id"]) == "" {
|
||||
return true, fmt.Errorf("missing reasoning item id")
|
||||
}
|
||||
if s.reasoning == nil {
|
||||
s.reasoning = map[int]*reasoningItem{}
|
||||
}
|
||||
s.reasoning[key] = &reasoningItem{id: str(x["id"]), parts: map[string]bool{}}
|
||||
return true, nil
|
||||
}
|
||||
if r == nil {
|
||||
if known || str(x["type"]) == "reasoning" {
|
||||
return true, fmt.Errorf("reasoning event outside reasoning item")
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
if r.closed {
|
||||
return true, fmt.Errorf("event after reasoning item closed")
|
||||
}
|
||||
if typ == "response.output_item.done" {
|
||||
if str(x["type"]) != "reasoning" || str(x["id"]) != r.id {
|
||||
return true, fmt.Errorf("reasoning item identity mismatch")
|
||||
}
|
||||
if err := validateReasoning(x); err != nil {
|
||||
return true, err
|
||||
}
|
||||
for _, closed := range r.parts {
|
||||
if !closed {
|
||||
return true, fmt.Errorf("reasoning item has open parts")
|
||||
}
|
||||
}
|
||||
r.closed = true
|
||||
return true, nil
|
||||
}
|
||||
if str(m["item_id"]) != r.id {
|
||||
return true, fmt.Errorf("reasoning event item identity mismatch")
|
||||
}
|
||||
field, partType := "content_index", "reasoning_text"
|
||||
switch typ {
|
||||
case "response.reasoning_summary_text.delta", "response.reasoning_summary_text.done",
|
||||
"response.reasoning_summary_part.added", "response.reasoning_summary_part.done":
|
||||
field, partType = "summary_index", "summary_text"
|
||||
case "response.reasoning_text.delta", "response.reasoning_text.done",
|
||||
"response.content_part.added", "response.content_part.done":
|
||||
default:
|
||||
return true, fmt.Errorf("non-reasoning event on reasoning item")
|
||||
}
|
||||
i, err := responseIndex(m, field)
|
||||
if err != nil {
|
||||
return true, err
|
||||
}
|
||||
partKey := fmt.Sprintf("%s:%d", field, i)
|
||||
closed, exists := r.parts[partKey]
|
||||
switch typ {
|
||||
case "response.content_part.added", "response.reasoning_summary_part.added",
|
||||
"response.content_part.done", "response.reasoning_summary_part.done":
|
||||
p := obj(m["part"])
|
||||
if str(p["type"]) != partType {
|
||||
return true, fmt.Errorf("unsupported reasoning part")
|
||||
}
|
||||
if err := fields(p, "type", "text"); err != nil {
|
||||
return true, err
|
||||
}
|
||||
added := typ == "response.content_part.added" || typ == "response.reasoning_summary_part.added"
|
||||
if (added && exists) || (!added && (!exists || closed)) {
|
||||
return true, fmt.Errorf("invalid reasoning part lifecycle")
|
||||
}
|
||||
r.parts[partKey] = !added
|
||||
default:
|
||||
if !exists || closed {
|
||||
return true, fmt.Errorf("reasoning text outside active part")
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const reasoningJSON = `{"type":"reasoning","id":"r","content":[{"type":"reasoning_text","text":"SECRET"}],"summary":[{"type":"summary_text","text":"SUMMARY"}],"encrypted_content":"CIPHER"}`
|
||||
const visibleJSON = `{"type":"message","id":"m","role":"assistant","phase":"final_answer","content":[{"type":"output_text","text":"OK","logprobs":[]}]}`
|
||||
|
||||
func reasoningSSE(events ...string) string {
|
||||
sep := string([]byte{10, 10})
|
||||
return "data: " + strings.Join(events, sep+"data: ") + sep
|
||||
}
|
||||
func TestReasoningConversions(t *testing.T) {
|
||||
for _, tool := range []bool{false, true} {
|
||||
for _, incomplete := range []bool{false, true} {
|
||||
visible, initial, delta := visibleJSON, `{"type":"message","content":[]}`, `{"type":"response.output_text.delta","output_index":1,"delta":"OK"}`
|
||||
if tool {
|
||||
visible = `{"type":"function_call","call_id":"c","name":"lookup","arguments":"{}"}`
|
||||
initial = `{"type":"function_call","call_id":"c","name":"lookup","arguments":""}`
|
||||
delta = `{"type":"response.function_call_arguments.delta","output_index":1,"delta":"{}"}`
|
||||
}
|
||||
status, details := "completed", ""
|
||||
if incomplete {
|
||||
status = "incomplete"
|
||||
details = `,"incomplete_details":{"reason":"max_output_tokens"}`
|
||||
}
|
||||
body := `{"id":"res","model":"m","status":"` + status + `","output":[` + reasoningJSON + `,` + visible + `],"usage":{"input_tokens":9,"output_tokens":7,"output_tokens_details":{"reasoning_tokens":5}}` + details + `}`
|
||||
src := reasoningSSE(`{"type":"response.output_item.added","output_index":0,"item":{"type":"reasoning","id":"r"}}`,
|
||||
`{"type":"response.content_part.added","output_index":0,"item_id":"r","content_index":0,"part":{"type":"reasoning_text","text":""}}`,
|
||||
`{"type":"response.reasoning_text.delta","output_index":0,"item_id":"r","content_index":0,"delta":"SECRET"}`,
|
||||
`{"type":"response.reasoning_text.done","output_index":0,"item_id":"r","content_index":0,"text":"SECRET"}`,
|
||||
`{"type":"response.content_part.done","output_index":0,"item_id":"r","content_index":0,"part":{"type":"reasoning_text","text":"SECRET"}}`,
|
||||
`{"type":"response.reasoning_summary_part.added","output_index":0,"item_id":"r","summary_index":0,"part":{"type":"summary_text","text":""}}`,
|
||||
`{"type":"response.reasoning_summary_text.delta","output_index":0,"item_id":"r","summary_index":0,"delta":"SUMMARY"}`,
|
||||
`{"type":"response.reasoning_summary_text.done","output_index":0,"item_id":"r","summary_index":0,"text":"SUMMARY"}`,
|
||||
`{"type":"response.reasoning_summary_part.done","output_index":0,"item_id":"r","summary_index":0,"part":{"type":"summary_text","text":"SUMMARY"}}`,
|
||||
`{"type":"response.output_item.done","output_index":0,"item":`+reasoningJSON+`}`,
|
||||
`{"type":"response.output_item.added","output_index":1,"item":`+initial+`}`, delta,
|
||||
`{"type":"response.output_item.done","output_index":1,"item":`+visible+`}`,
|
||||
`{"type":"response.`+status+`","response":`+body+`}`)
|
||||
for _, target := range []Kind{Chat, Anthropic, Responses} {
|
||||
out, err := ConvertResponse(Responses, target, []byte(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var b bytes.Buffer
|
||||
u, err := Stream(context.Background(), &b, strings.NewReader(src), Responses, target, "m")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !u.Known || u.Input != 9 || u.Output != 7 || !bytes.Contains(u.Raw, []byte("reasoning_tokens")) {
|
||||
t.Fatalf("usage: %+v", u)
|
||||
}
|
||||
if target == Responses {
|
||||
if string(out) != body || b.String() != src {
|
||||
t.Fatal("changed same-protocol bytes")
|
||||
}
|
||||
continue
|
||||
}
|
||||
for _, result := range []string{string(out), b.String()} {
|
||||
for _, secret := range []string{"SECRET", "SUMMARY", "CIPHER", "signature", "thinking", "phase"} {
|
||||
if strings.Contains(result, secret) {
|
||||
t.Fatal("leaked " + secret)
|
||||
}
|
||||
}
|
||||
want := "OK"
|
||||
if tool {
|
||||
want = "lookup"
|
||||
}
|
||||
if !strings.Contains(result, want) {
|
||||
t.Fatal("lost visible content")
|
||||
}
|
||||
finish := "stop"
|
||||
if target == Anthropic {
|
||||
finish = "end_turn"
|
||||
}
|
||||
if tool {
|
||||
finish = "tool_calls"
|
||||
if target == Anthropic {
|
||||
finish = "tool_use"
|
||||
}
|
||||
}
|
||||
if incomplete {
|
||||
finish = "length"
|
||||
if target == Anthropic {
|
||||
finish = "max_tokens"
|
||||
}
|
||||
}
|
||||
if !strings.Contains(result, finish) {
|
||||
t.Fatal("lost finish " + finish)
|
||||
}
|
||||
}
|
||||
if target == Chat && !strings.Contains(b.String(), "[DONE]") {
|
||||
t.Fatal("missing terminal")
|
||||
}
|
||||
if target == Anthropic && (!strings.Contains(b.String(), "message_stop") || !strings.Contains(b.String(), `"index":0`)) {
|
||||
t.Fatal("terminal/index")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReasoningRejectsInvalidStreams(t *testing.T) {
|
||||
start := `{"type":"response.output_item.added","output_index":0,"item":{"type":"reasoning","id":"r"}}`
|
||||
part := `{"type":"response.reasoning_summary_part.added","output_index":0,"item_id":"r","summary_index":0,"part":{"type":"summary_text","text":""}}`
|
||||
done := `{"type":"response.output_item.done","output_index":0,"item":` + reasoningJSON + `}`
|
||||
for name, events := range map[string][]string{
|
||||
"duplicate": {start, start}, "part before item": {part}, "open part": {start, part, done}, "closed item": {start, done, part},
|
||||
"wrong id": {start, strings.Replace(part, `"item_id":"r"`, `"item_id":"x"`, 1)},
|
||||
"unknown": {start, `{"type":"response.reasoning_unknown.delta","output_index":0}`},
|
||||
"text collision": {start, `{"type":"response.output_text.delta","output_index":0,"item_id":"r","delta":"lost"}`},
|
||||
"tool collision": {start, `{"type":"response.function_call_arguments.delta","output_index":0,"item_id":"r","delta":"{}"}`},
|
||||
"unknown part": {start, strings.Replace(part, "summary_text", "output_text", 1)},
|
||||
"negative": {strings.Replace(start, `"output_index":0`, `"output_index":-1`, 1)},
|
||||
"fractional": {strings.Replace(start, `"output_index":0`, `"output_index":0.5`, 1)},
|
||||
"final missing": {start, done, `{"type":"response.completed","response":{"status":"completed","output":[]}}`},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
var b bytes.Buffer
|
||||
_, err := Stream(context.Background(), &b, strings.NewReader(reasoningSSE(events...)), Responses, Chat, "")
|
||||
if err == nil || strings.Contains(err.Error(), "truncat") {
|
||||
t.Fatalf("expected validation error: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, bad := range []string{strings.Replace(reasoningJSON, "reasoning_text", "output_text", 1), strings.Replace(reasoningJSON, "summary_text", "unknown", 1), `{"type":"reasoning","summary":{}}`, strings.Replace(visibleJSON, `"logprobs":[]`, `"logprobs":[{}]`, 1), strings.Replace(visibleJSON, `"phase":`, `"unknown":`, 1)} {
|
||||
if _, err := ConvertResponse(Responses, Chat, []byte(`{"status":"completed","output":[`+bad+`]}`)); err == nil {
|
||||
t.Fatal("accepted invalid output")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,879 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const maxStreamBytes = 16 << 20
|
||||
const maxStreamBlocks = 4096
|
||||
const maxStreamUsageBytes = 16 << 10
|
||||
const maxPassthroughBytes = 64 << 20
|
||||
|
||||
type streamBlock struct {
|
||||
part
|
||||
content, identity, functionName strings.Builder
|
||||
index int
|
||||
closed bool
|
||||
emitted bool
|
||||
sent int
|
||||
}
|
||||
type streamState struct {
|
||||
thinking map[int]*thinkingBlock
|
||||
ctx context.Context
|
||||
w io.Writer
|
||||
from, to Kind
|
||||
r response
|
||||
blocks []*streamBlock
|
||||
source map[int]*streamBlock
|
||||
reasoning map[int]*reasoningItem
|
||||
started, finished, terminal bool
|
||||
sequence int
|
||||
retained, blockCount int
|
||||
byteLimit, blockLimit int
|
||||
}
|
||||
|
||||
func (s *streamState) retain(n int) error {
|
||||
if n > s.byteLimit-s.retained {
|
||||
return fmt.Errorf("stream retained content exceeds %d bytes", s.byteLimit)
|
||||
}
|
||||
s.retained += n
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) reserveBlock() error {
|
||||
if s.blockCount >= s.blockLimit {
|
||||
return fmt.Errorf("stream exceeds %d blocks", s.blockLimit)
|
||||
}
|
||||
s.blockCount++
|
||||
return nil
|
||||
}
|
||||
|
||||
func writeFlush(w io.Writer, b []byte) error {
|
||||
n, e := w.Write(b)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
if n != len(b) {
|
||||
return io.ErrShortWrite
|
||||
}
|
||||
if f, ok := w.(interface{ Flush() error }); ok {
|
||||
return f.Flush()
|
||||
}
|
||||
if f, ok := w.(interface{ Flush() }); ok {
|
||||
f.Flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) event(name string, m object) error {
|
||||
if e := s.ctx.Err(); e != nil {
|
||||
return e
|
||||
}
|
||||
if s.to == Responses {
|
||||
m["sequence_number"] = s.sequence
|
||||
s.sequence++
|
||||
}
|
||||
b, e := json.Marshal(m)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
prefix := ""
|
||||
if name != "" {
|
||||
prefix = "event: " + name + "\n"
|
||||
}
|
||||
return writeFlush(s.w, []byte(prefix+"data: "+string(b)+"\n\n"))
|
||||
}
|
||||
func (s *streamState) typed(name string, m object) error { m["type"] = name; return s.event(name, m) }
|
||||
func (s *streamState) chat(delta object, finish any) error {
|
||||
return s.event("", object{"id": s.r.id, "object": "chat.completion.chunk", "created": 0, "model": s.r.model, "choices": []any{object{"index": 0, "delta": delta, "finish_reason": finish}}})
|
||||
}
|
||||
func (s *streamState) start() error {
|
||||
if s.started {
|
||||
return nil
|
||||
}
|
||||
s.started = true
|
||||
if s.r.id == "" {
|
||||
s.r.id = newID()
|
||||
}
|
||||
switch s.to {
|
||||
case Chat:
|
||||
return s.chat(object{"role": "assistant", "content": ""}, nil)
|
||||
case Anthropic:
|
||||
m := object{"id": s.r.id, "type": "message", "role": "assistant", "model": s.r.model, "content": []any{}, "stop_reason": nil, "stop_sequence": nil}
|
||||
m["usage"] = usageJSON(Anthropic, s.r.usage)
|
||||
return s.typed("message_start", object{"message": m})
|
||||
case Responses:
|
||||
r := object{"id": s.r.id, "object": "response", "created_at": time.Now().Unix(), "model": s.r.model, "status": "in_progress", "output": []any{}, "error": nil, "incomplete_details": nil, "usage": nil}
|
||||
if e := s.typed("response.created", object{"response": r}); e != nil {
|
||||
return e
|
||||
}
|
||||
return s.typed("response.in_progress", object{"response": r})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) itemID(b *streamBlock) string {
|
||||
if b.typ == "call" {
|
||||
return fmt.Sprintf("fc_%s_%d", s.r.id, b.index)
|
||||
}
|
||||
return fmt.Sprintf("msg_%s_%d", s.r.id, b.index)
|
||||
}
|
||||
func (s *streamState) item(b *streamBlock, done bool) object {
|
||||
status := "in_progress"
|
||||
if done {
|
||||
status = "completed"
|
||||
}
|
||||
if b.typ == "call" {
|
||||
args := ""
|
||||
if done {
|
||||
args = b.args
|
||||
}
|
||||
return object{"id": s.itemID(b), "type": "function_call", "call_id": b.id, "name": b.name, "arguments": args, "status": status}
|
||||
}
|
||||
content := []any{}
|
||||
if done {
|
||||
content = append(content, object{"type": "output_text", "text": b.text, "annotations": []any{}})
|
||||
}
|
||||
return object{"id": s.itemID(b), "type": "message", "role": "assistant", "status": status, "content": content}
|
||||
}
|
||||
func (s *streamState) open(b *streamBlock) error {
|
||||
if b.emitted {
|
||||
return nil
|
||||
}
|
||||
if b.typ == "call" && (b.id == "" || b.name == "") {
|
||||
return nil
|
||||
}
|
||||
if e := s.start(); e != nil {
|
||||
return e
|
||||
}
|
||||
b.emitted = true
|
||||
switch s.to {
|
||||
case Chat:
|
||||
if b.typ == "call" {
|
||||
return s.chat(object{"tool_calls": []any{object{"index": s.toolIndex(b), "id": b.id, "type": "function", "function": object{"name": b.name, "arguments": ""}}}}, nil)
|
||||
}
|
||||
case Anthropic:
|
||||
c := object{"type": "text", "text": ""}
|
||||
if b.typ == "call" {
|
||||
c = object{"type": "tool_use", "id": b.id, "name": b.name, "input": object{}}
|
||||
}
|
||||
return s.typed("content_block_start", object{"index": b.index, "content_block": c})
|
||||
case Responses:
|
||||
if e := s.typed("response.output_item.added", object{"output_index": b.index, "item": s.item(b, false)}); e != nil {
|
||||
return e
|
||||
}
|
||||
if b.typ == "text" {
|
||||
return s.typed("response.content_part.added", object{"output_index": b.index, "item_id": s.itemID(b), "content_index": 0, "part": object{"type": "output_text", "text": "", "annotations": []any{}}})
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) toolIndex(b *streamBlock) int {
|
||||
i := 0
|
||||
for _, x := range s.blocks {
|
||||
if x == b {
|
||||
return i
|
||||
}
|
||||
if x.typ == "call" {
|
||||
i++
|
||||
}
|
||||
}
|
||||
return i
|
||||
}
|
||||
func (s *streamState) add(key int, p part) (*streamBlock, error) {
|
||||
if s.finished {
|
||||
return nil, fmt.Errorf("content after finish")
|
||||
}
|
||||
if _, ok := s.source[key]; ok {
|
||||
return nil, fmt.Errorf("duplicate content index %d", key)
|
||||
}
|
||||
if e := s.reserveBlock(); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if e := s.retain(len(p.text) + len(p.args) + len(p.id) + len(p.name)); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
b := &streamBlock{part: p, index: len(s.blocks)}
|
||||
b.content.WriteString(p.text)
|
||||
b.content.WriteString(p.args)
|
||||
b.identity.WriteString(p.id)
|
||||
b.functionName.WriteString(p.name)
|
||||
s.blocks = append(s.blocks, b)
|
||||
s.source[key] = b
|
||||
return b, s.open(b)
|
||||
}
|
||||
func (s *streamState) delta(b *streamBlock, text string) error {
|
||||
if b == nil || b.closed || s.finished {
|
||||
return fmt.Errorf("delta outside active content block")
|
||||
}
|
||||
if e := s.retain(len(text)); e != nil {
|
||||
return e
|
||||
}
|
||||
b.content.WriteString(text)
|
||||
// Builder.String shares storage without copying accumulated content.
|
||||
if b.typ == "text" {
|
||||
b.text = b.content.String()
|
||||
} else {
|
||||
b.args = b.content.String()
|
||||
}
|
||||
if e := s.open(b); e != nil {
|
||||
return e
|
||||
}
|
||||
if !b.emitted {
|
||||
return nil
|
||||
}
|
||||
whole := b.text
|
||||
if b.typ == "call" {
|
||||
whole = b.args
|
||||
}
|
||||
text = whole[b.sent:]
|
||||
b.sent = len(whole)
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
switch s.to {
|
||||
case Chat:
|
||||
if b.typ == "text" {
|
||||
return s.chat(object{"content": text}, nil)
|
||||
}
|
||||
return s.chat(object{"tool_calls": []any{object{"index": s.toolIndex(b), "function": object{"arguments": text}}}}, nil)
|
||||
case Anthropic:
|
||||
d := object{"type": "text_delta", "text": text}
|
||||
if b.typ == "call" {
|
||||
d = object{"type": "input_json_delta", "partial_json": text}
|
||||
}
|
||||
return s.typed("content_block_delta", object{"index": b.index, "delta": d})
|
||||
case Responses:
|
||||
name := "response.output_text.delta"
|
||||
m := object{"output_index": b.index, "item_id": s.itemID(b), "delta": text}
|
||||
if b.typ == "call" {
|
||||
name = "response.function_call_arguments.delta"
|
||||
} else {
|
||||
m["content_index"] = 0
|
||||
}
|
||||
return s.typed(name, m)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) close(b *streamBlock) error {
|
||||
if b == nil {
|
||||
return fmt.Errorf("unknown content block")
|
||||
}
|
||||
if b.closed {
|
||||
return nil
|
||||
}
|
||||
if b.typ == "call" {
|
||||
if b.id == "" || b.name == "" {
|
||||
return fmt.Errorf("incomplete tool identity")
|
||||
}
|
||||
if b.args == "" {
|
||||
if e := s.delta(b, "{}"); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
if _, e := decode([]byte(b.args)); e != nil {
|
||||
return fmt.Errorf("invalid streamed tool arguments: %w", e)
|
||||
}
|
||||
}
|
||||
if e := s.open(b); e != nil {
|
||||
return e
|
||||
}
|
||||
if e := s.delta(b, ""); e != nil {
|
||||
return e
|
||||
}
|
||||
b.closed = true
|
||||
switch s.to {
|
||||
case Anthropic:
|
||||
return s.typed("content_block_stop", object{"index": b.index})
|
||||
case Responses:
|
||||
m := object{"output_index": b.index, "item_id": s.itemID(b)}
|
||||
name := "response.function_call_arguments.done"
|
||||
if b.typ == "call" {
|
||||
m["arguments"] = b.args
|
||||
} else {
|
||||
name = "response.output_text.done"
|
||||
m["text"] = b.text
|
||||
m["content_index"] = 0
|
||||
}
|
||||
if e := s.typed(name, m); e != nil {
|
||||
return e
|
||||
}
|
||||
if b.typ == "text" {
|
||||
if e := s.typed("response.content_part.done", object{"output_index": b.index, "item_id": s.itemID(b), "content_index": 0, "part": object{"type": "output_text", "text": b.text, "annotations": []any{}}}); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
return s.typed("response.output_item.done", object{"output_index": b.index, "item": s.item(b, true)})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) finish(reason string) error {
|
||||
if s.finished {
|
||||
return fmt.Errorf("duplicate finish")
|
||||
}
|
||||
if e := s.start(); e != nil {
|
||||
return e
|
||||
}
|
||||
for _, b := range s.blocks {
|
||||
if e := s.close(b); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
s.r.finish = reason
|
||||
s.finished = true
|
||||
return nil
|
||||
}
|
||||
func (s *streamState) end() error {
|
||||
if !s.finished {
|
||||
return fmt.Errorf("stream ended without finish reason")
|
||||
}
|
||||
switch s.to {
|
||||
case Chat:
|
||||
if e := s.chat(object{}, wireFinish(Chat, s.r.finish)); e != nil {
|
||||
return e
|
||||
}
|
||||
if s.r.usage.Known {
|
||||
if e := s.event("", object{"id": s.r.id, "object": "chat.completion.chunk", "created": 0, "model": s.r.model, "choices": []any{}, "usage": usageJSON(Chat, s.r.usage)}); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
return writeFlush(s.w, []byte("data: [DONE]\n\n"))
|
||||
case Anthropic:
|
||||
m := object{"delta": object{"stop_reason": wireFinish(Anthropic, s.r.finish), "stop_sequence": nil}}
|
||||
if s.r.usage.Known {
|
||||
m["usage"] = usageJSON(Anthropic, s.r.usage)
|
||||
}
|
||||
if e := s.typed("message_delta", m); e != nil {
|
||||
return e
|
||||
}
|
||||
return s.typed("message_stop", object{})
|
||||
case Responses:
|
||||
r := s.r
|
||||
for _, b := range s.blocks {
|
||||
r.parts = append(r.parts, b.part)
|
||||
}
|
||||
m, e := writeResponse(Responses, r)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
items := []any{}
|
||||
for _, b := range s.blocks {
|
||||
items = append(items, s.item(b, true))
|
||||
}
|
||||
m["output"] = items
|
||||
name := "response.completed"
|
||||
if r.finish == "length" {
|
||||
name = "response.incomplete"
|
||||
}
|
||||
return s.typed(name, object{"response": m})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func mergeUsage(k Kind, old Usage, m object) (Usage, error) {
|
||||
if r := obj(m["response"]); r != nil {
|
||||
m = r
|
||||
}
|
||||
if r := obj(m["message"]); r != nil {
|
||||
m = r
|
||||
}
|
||||
u := obj(m["usage"])
|
||||
if u == nil {
|
||||
return old, nil
|
||||
}
|
||||
raw, err := json.Marshal(u)
|
||||
if err != nil {
|
||||
return old, err
|
||||
}
|
||||
if len(raw) > maxStreamUsageBytes {
|
||||
return old, fmt.Errorf("stream usage exceeds %d bytes", maxStreamUsageBytes)
|
||||
}
|
||||
if len(old.Raw) > 0 {
|
||||
prior, e := decode(old.Raw)
|
||||
if e != nil {
|
||||
return old, e
|
||||
}
|
||||
for key, v := range u {
|
||||
prior[key] = v
|
||||
}
|
||||
u = prior
|
||||
}
|
||||
n := usage(k, u)
|
||||
if len(n.Raw) > maxStreamUsageBytes {
|
||||
return old, fmt.Errorf("stream usage exceeds %d bytes", maxStreamUsageBytes)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
func streamError(m object) error {
|
||||
if m["error"] != nil {
|
||||
return fmt.Errorf("upstream stream error: %v", m["error"])
|
||||
}
|
||||
switch str(m["type"]) {
|
||||
case "error", "response.failed":
|
||||
return fmt.Errorf("upstream stream error: %v", m)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *streamState) consume(m object, event string) error {
|
||||
if e := streamError(m); e != nil {
|
||||
return e
|
||||
}
|
||||
u, err := mergeUsage(s.from, s.r.usage, m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.r.usage = u
|
||||
if s.from == Chat {
|
||||
if id := str(m["id"]); id != "" && !s.started {
|
||||
s.r.id = id
|
||||
}
|
||||
if s.r.model == "" {
|
||||
s.r.model = str(m["model"])
|
||||
}
|
||||
cs := arr(m["choices"])
|
||||
if len(cs) > 1 {
|
||||
return fmt.Errorf("multiple streaming choices unsupported")
|
||||
}
|
||||
for _, v := range cs {
|
||||
c := obj(v)
|
||||
if num(c["index"]) != 0 {
|
||||
return fmt.Errorf("nonzero choice index unsupported")
|
||||
}
|
||||
if c["logprobs"] != nil {
|
||||
return fmt.Errorf("logprobs unsupported")
|
||||
}
|
||||
d, e := chatOutput(obj(c["delta"]))
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
if e := fields(d, "role", "content", "tool_calls"); e != nil {
|
||||
return e
|
||||
}
|
||||
if role := str(d["role"]); role != "" && role != "assistant" {
|
||||
return fmt.Errorf("invalid delta role")
|
||||
}
|
||||
if e := s.start(); e != nil {
|
||||
return e
|
||||
}
|
||||
if text, ok := d["content"]; ok && text != nil {
|
||||
t, ok := text.(string)
|
||||
if !ok {
|
||||
return fmt.Errorf("non-text streamed content")
|
||||
}
|
||||
b := s.source[-1]
|
||||
if b == nil {
|
||||
var e error
|
||||
b, e = s.add(-1, part{typ: "text"})
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
if e := s.delta(b, t); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
for _, v := range arr(d["tool_calls"]) {
|
||||
t := obj(v)
|
||||
if e := fields(t, "index", "id", "type", "function"); e != nil {
|
||||
return e
|
||||
}
|
||||
idx, ok := t["index"]
|
||||
if !ok || num(idx) < 0 {
|
||||
return fmt.Errorf("invalid tool index")
|
||||
}
|
||||
if typ := str(t["type"]); typ != "" && typ != "function" {
|
||||
return fmt.Errorf("unsupported streamed tool")
|
||||
}
|
||||
f := obj(t["function"])
|
||||
if e := fields(f, "name", "arguments"); e != nil {
|
||||
return e
|
||||
}
|
||||
key := int(num(idx))
|
||||
b := s.source[key]
|
||||
if b == nil {
|
||||
var e error
|
||||
b, e = s.add(key, part{typ: "call"})
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
if b.emitted && (str(t["id"]) != "" || str(f["name"]) != "") {
|
||||
return fmt.Errorf("fragmented tool identity after arguments unsupported")
|
||||
}
|
||||
if e := s.retain(len(str(t["id"])) + len(str(f["name"]))); e != nil {
|
||||
return e
|
||||
}
|
||||
b.identity.WriteString(str(t["id"]))
|
||||
b.functionName.WriteString(str(f["name"]))
|
||||
b.id = b.identity.String()
|
||||
b.name = b.functionName.String()
|
||||
if e := s.delta(b, str(f["arguments"])); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
if reason := str(c["finish_reason"]); reason != "" {
|
||||
f, e := finishReason(Chat, reason)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
if e = s.finish(f); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
typ := str(m["type"])
|
||||
if typ == "" {
|
||||
typ = event
|
||||
}
|
||||
if event != "" && typ != event {
|
||||
return fmt.Errorf("SSE event/type mismatch")
|
||||
}
|
||||
// Discarded blocks still retain lifecycle bookkeeping in separate maps.
|
||||
if (s.from == Anthropic && typ == "content_block_start" && (str(obj(m["content_block"])["type"]) == "thinking" || str(obj(m["content_block"])["type"]) == "redacted_thinking")) ||
|
||||
(s.from == Responses && typ == "response.output_item.added" && str(obj(m["item"])["type"]) == "reasoning") ||
|
||||
(s.from == Responses && (typ == "response.content_part.added" || typ == "response.reasoning_summary_part.added") && s.reasoning[int(num(m["output_index"]))] != nil) {
|
||||
if e := s.reserveBlock(); e != nil {
|
||||
return e
|
||||
}
|
||||
if typ == "response.output_item.added" {
|
||||
if e := s.retain(len(str(obj(m["item"])["id"]))); e != nil {
|
||||
return e
|
||||
}
|
||||
}
|
||||
}
|
||||
if s.from == Anthropic {
|
||||
if handled, err := s.thinkingEvent(typ, m); handled || err != nil {
|
||||
return err
|
||||
}
|
||||
switch typ {
|
||||
case "ping":
|
||||
return nil
|
||||
case "message_start":
|
||||
if s.started {
|
||||
return fmt.Errorf("duplicate message_start")
|
||||
}
|
||||
x := obj(m["message"])
|
||||
s.r.id = str(x["id"])
|
||||
if s.r.model == "" {
|
||||
s.r.model = str(x["model"])
|
||||
}
|
||||
if len(arr(x["content"])) != 0 {
|
||||
return fmt.Errorf("nonempty message_start content")
|
||||
}
|
||||
return s.start()
|
||||
case "content_block_start":
|
||||
x := obj(m["content_block"])
|
||||
p := part{typ: "text"}
|
||||
initial := ""
|
||||
switch str(x["type"]) {
|
||||
case "text":
|
||||
if e := fields(x, "type", "text"); e != nil {
|
||||
return e
|
||||
}
|
||||
initial = str(x["text"])
|
||||
case "tool_use":
|
||||
if e := fields(x, "type", "id", "name", "input"); e != nil {
|
||||
return e
|
||||
}
|
||||
p = part{typ: "call", id: str(x["id"]), name: str(x["name"])}
|
||||
if obj(x["input"]) == nil {
|
||||
return fmt.Errorf("tool input must be an object")
|
||||
}
|
||||
if len(obj(x["input"])) > 0 {
|
||||
b, e := json.Marshal(x["input"])
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
initial = string(b)
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("unsupported streamed block %q", str(x["type"]))
|
||||
}
|
||||
b, e := s.add(int(num(m["index"])), p)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
return s.delta(b, initial)
|
||||
case "content_block_delta":
|
||||
b := s.source[int(num(m["index"]))]
|
||||
d := obj(m["delta"])
|
||||
switch str(d["type"]) {
|
||||
case "text_delta":
|
||||
if b == nil || b.typ != "text" {
|
||||
return fmt.Errorf("text delta block mismatch")
|
||||
}
|
||||
return s.delta(b, str(d["text"]))
|
||||
case "input_json_delta":
|
||||
if b == nil || b.typ != "call" {
|
||||
return fmt.Errorf("tool delta block mismatch")
|
||||
}
|
||||
return s.delta(b, str(d["partial_json"]))
|
||||
default:
|
||||
return fmt.Errorf("unsupported delta %q", str(d["type"]))
|
||||
}
|
||||
case "content_block_stop":
|
||||
return s.close(s.source[int(num(m["index"]))])
|
||||
case "message_delta":
|
||||
d := obj(m["delta"])
|
||||
if e := fields(d, "stop_reason", "stop_sequence"); e != nil {
|
||||
return e
|
||||
}
|
||||
if d["stop_sequence"] != nil {
|
||||
return fmt.Errorf("stop sequence cannot be preserved")
|
||||
}
|
||||
f, e := finishReason(Anthropic, str(d["stop_reason"]))
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
return s.finish(f)
|
||||
case "message_stop":
|
||||
s.terminal = true
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("unsupported Anthropic event %q", typ)
|
||||
}
|
||||
}
|
||||
if handled, err := s.reasoningEvent(typ, m); handled || err != nil {
|
||||
return err
|
||||
}
|
||||
switch typ {
|
||||
case "response.created", "response.in_progress":
|
||||
x := obj(m["response"])
|
||||
if !s.started {
|
||||
s.r.id = str(x["id"])
|
||||
if s.r.model == "" {
|
||||
s.r.model = str(x["model"])
|
||||
}
|
||||
}
|
||||
return s.start()
|
||||
case "response.output_item.added":
|
||||
x := obj(m["item"])
|
||||
key := int(num(m["output_index"]))
|
||||
p := part{typ: "text"}
|
||||
switch str(x["type"]) {
|
||||
case "message":
|
||||
if len(arr(x["content"])) > 0 {
|
||||
return fmt.Errorf("nonempty initial response content")
|
||||
}
|
||||
case "function_call":
|
||||
if _, ok := x["arguments"].(string); !ok {
|
||||
return fmt.Errorf("function arguments must be a string")
|
||||
}
|
||||
p = part{typ: "call", id: str(x["call_id"]), name: str(x["name"])}
|
||||
default:
|
||||
return fmt.Errorf("unsupported output item %q", str(x["type"]))
|
||||
}
|
||||
b, e := s.add(key, p)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
if p.typ == "call" {
|
||||
return s.delta(b, str(x["arguments"]))
|
||||
}
|
||||
return nil
|
||||
case "response.content_part.added":
|
||||
p := obj(m["part"])
|
||||
if str(p["type"]) != "output_text" || num(m["content_index"]) != 0 {
|
||||
return fmt.Errorf("unsupported response content part")
|
||||
}
|
||||
if len(arr(p["annotations"])) > 0 {
|
||||
return fmt.Errorf("annotations unsupported")
|
||||
}
|
||||
return s.delta(s.source[int(num(m["output_index"]))], str(p["text"]))
|
||||
case "response.output_text.delta", "response.function_call_arguments.delta":
|
||||
b := s.source[int(num(m["output_index"]))]
|
||||
if b == nil {
|
||||
return fmt.Errorf("delta before output item")
|
||||
}
|
||||
if (typ == "response.output_text.delta") != (b.typ == "text") || num(m["content_index"]) != 0 {
|
||||
return fmt.Errorf("response delta block mismatch")
|
||||
}
|
||||
return s.delta(b, str(m["delta"]))
|
||||
case "response.output_text.done", "response.function_call_arguments.done":
|
||||
b := s.source[int(num(m["output_index"]))]
|
||||
if b == nil {
|
||||
return fmt.Errorf("done before output item")
|
||||
}
|
||||
full := str(m["text"])
|
||||
got := b.text
|
||||
if b.typ == "call" {
|
||||
full = str(m["arguments"])
|
||||
got = b.args
|
||||
}
|
||||
if full != got {
|
||||
return fmt.Errorf("final content disagrees with deltas")
|
||||
}
|
||||
return nil
|
||||
case "response.content_part.done":
|
||||
p := obj(m["part"])
|
||||
if str(p["type"]) != "output_text" || len(arr(p["annotations"])) > 0 || num(m["content_index"]) != 0 {
|
||||
return fmt.Errorf("unsupported content part")
|
||||
}
|
||||
return nil
|
||||
case "response.output_item.done":
|
||||
return s.close(s.source[int(num(m["output_index"]))])
|
||||
case "response.completed", "response.incomplete":
|
||||
r, e := readResponse(Responses, obj(m["response"]))
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
if len(r.parts) != len(s.blocks) {
|
||||
return fmt.Errorf("final response disagrees with streamed items")
|
||||
}
|
||||
for i, p := range r.parts {
|
||||
b := s.blocks[i]
|
||||
if p.typ != b.typ || p.text != b.text || p.args != b.args || p.id != b.id || p.name != b.name {
|
||||
return fmt.Errorf("final response disagrees with streamed content")
|
||||
}
|
||||
}
|
||||
if e = s.finish(r.finish); e != nil {
|
||||
return e
|
||||
}
|
||||
s.terminal = true
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("unsupported Responses event %q", typ)
|
||||
}
|
||||
}
|
||||
|
||||
// Stream does not own r. A caller must cancel or close a blocked upstream read.
|
||||
func Stream(ctx context.Context, w io.Writer, r io.Reader, from, to Kind, model string) (Usage, error) {
|
||||
return streamWithLimits(ctx, w, r, from, to, model, maxStreamBytes, maxStreamBlocks, maxPassthroughBytes)
|
||||
}
|
||||
|
||||
func readSSELine(ctx context.Context, br *bufio.Reader, remaining int) (string, error) {
|
||||
var line strings.Builder
|
||||
for {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
fragment, err := br.ReadSlice('\n')
|
||||
if len(fragment) > remaining-line.Len() {
|
||||
return "", fmt.Errorf("SSE line/event exceeds size limit")
|
||||
}
|
||||
line.Write(fragment)
|
||||
if err != bufio.ErrBufferFull {
|
||||
return line.String(), err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func streamWithLimits(ctx context.Context, w io.Writer, r io.Reader, from, to Kind, model string, byteLimit, blockLimit, totalLimit int) (Usage, error) {
|
||||
if !valid(from) || !valid(to) {
|
||||
return Usage{}, fmt.Errorf("unknown protocol")
|
||||
}
|
||||
s := &streamState{ctx: ctx, w: w, from: from, to: to, r: response{model: model}, source: map[int]*streamBlock{}, byteLimit: byteLimit, blockLimit: blockLimit}
|
||||
br := bufio.NewReader(r)
|
||||
data := []string{}
|
||||
event := ""
|
||||
size := 0
|
||||
total := 0
|
||||
dispatch := func() error {
|
||||
if len(data) == 0 {
|
||||
event = ""
|
||||
return nil
|
||||
}
|
||||
payload := strings.Join(data, "\n")
|
||||
data = nil
|
||||
ev := event
|
||||
event = ""
|
||||
size = 0
|
||||
if payload == "[DONE]" {
|
||||
if from != Chat {
|
||||
return fmt.Errorf("unexpected DONE marker")
|
||||
}
|
||||
s.terminal = true
|
||||
return nil
|
||||
}
|
||||
m, e := decode([]byte(payload))
|
||||
if e != nil {
|
||||
return fmt.Errorf("invalid SSE JSON: %w", e)
|
||||
}
|
||||
if e = streamError(m); e != nil {
|
||||
return e
|
||||
}
|
||||
if from == to {
|
||||
u, err := mergeUsage(from, s.r.usage, m)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.r.usage = u
|
||||
switch str(m["type"]) {
|
||||
case "message_stop", "response.completed", "response.incomplete":
|
||||
s.terminal = true
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if s.terminal {
|
||||
return fmt.Errorf("data after terminal event")
|
||||
}
|
||||
return s.consume(m, ev)
|
||||
}
|
||||
for {
|
||||
if e := ctx.Err(); e != nil {
|
||||
return s.r.usage, e
|
||||
}
|
||||
line, e := readSSELine(ctx, br, byteLimit-size)
|
||||
if e != nil && e != io.EOF {
|
||||
return s.r.usage, e
|
||||
}
|
||||
if len(line) > 0 {
|
||||
if from == to {
|
||||
if total+len(line) > totalLimit {
|
||||
return s.r.usage, fmt.Errorf("stream passthrough exceeds %d bytes", totalLimit)
|
||||
}
|
||||
total += len(line)
|
||||
if err := writeFlush(w, []byte(line)); err != nil {
|
||||
return s.r.usage, err
|
||||
}
|
||||
}
|
||||
size += len(line)
|
||||
line = strings.TrimSuffix(strings.TrimSuffix(line, "\n"), "\r")
|
||||
if line == "" {
|
||||
if err := dispatch(); err != nil {
|
||||
return s.r.usage, err
|
||||
}
|
||||
size = 0
|
||||
if s.terminal {
|
||||
break
|
||||
}
|
||||
} else if !strings.HasPrefix(line, ":") {
|
||||
key, val, _ := strings.Cut(line, ":")
|
||||
val = strings.TrimPrefix(val, " ")
|
||||
switch key {
|
||||
case "data":
|
||||
data = append(data, val)
|
||||
case "event":
|
||||
event = val
|
||||
}
|
||||
}
|
||||
}
|
||||
if e != nil {
|
||||
if e != io.EOF {
|
||||
return s.r.usage, e
|
||||
}
|
||||
if err := dispatch(); err != nil {
|
||||
return s.r.usage, err
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
if e := ctx.Err(); e != nil {
|
||||
return s.r.usage, e
|
||||
}
|
||||
if from == to {
|
||||
return s.r.usage, nil
|
||||
}
|
||||
if !s.terminal {
|
||||
return s.r.usage, fmt.Errorf("truncated SSE stream: missing terminal event")
|
||||
}
|
||||
if e := s.end(); e != nil {
|
||||
return s.r.usage, e
|
||||
}
|
||||
return s.r.usage, nil
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func limitEvent(s string) string { return "data: " + s + "\n\n" }
|
||||
func TestSSEReadLimits(t *testing.T) {
|
||||
for _, to := range []Kind{Chat, Anthropic} {
|
||||
var out bytes.Buffer
|
||||
_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat("x", 10000)), Chat, to, "", 64, 8, 64)
|
||||
if err == nil || !strings.Contains(err.Error(), "size limit") || out.Len() != 0 {
|
||||
t.Fatalf("%s: %v", to, err)
|
||||
}
|
||||
}
|
||||
for _, input := range []string{"12345678", "1234567\n"} {
|
||||
line, err := readSSELine(context.Background(), bufio.NewReader(strings.NewReader(input)), 8)
|
||||
if line != input || (err != nil && err != io.EOF) {
|
||||
t.Fatalf("%q %v", line, err)
|
||||
}
|
||||
}
|
||||
var out bytes.Buffer
|
||||
_, err := streamWithLimits(context.Background(), &out, strings.NewReader(":123456\n:123456\n"), Chat, Chat, "", 12, 8, 16)
|
||||
if err == nil || out.String() != ":123456\n" {
|
||||
t.Fatalf("%q %v", out.String(), err)
|
||||
}
|
||||
}
|
||||
func TestStreamPassthroughTotalLimit(t *testing.T) {
|
||||
frame := "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":null}]}\n\n"
|
||||
var out bytes.Buffer
|
||||
_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat(frame, 4)), Chat, Chat, "", 512, 8, 2*len(frame))
|
||||
if err == nil || !strings.Contains(err.Error(), "passthrough") || out.Len() != 2*len(frame) {
|
||||
t.Fatalf("out=%d %q %v", out.Len(), out.String(), err)
|
||||
}
|
||||
}
|
||||
func TestStreamCumulativeLimits(t *testing.T) {
|
||||
for _, field := range []string{"text", "arguments", "identity"} {
|
||||
delta := object{"content": strings.Repeat("x", 100)}
|
||||
if field != "text" {
|
||||
call := object{"index": 0}
|
||||
if field == "arguments" {
|
||||
call["function"] = object{"arguments": strings.Repeat("x", 100)}
|
||||
} else {
|
||||
call["id"] = strings.Repeat("x", 100)
|
||||
}
|
||||
delta = object{"tool_calls": []any{call}}
|
||||
}
|
||||
event := limitEvent(fmt.Sprintf("{\"choices\":[{\"index\":0,\"delta\":%s}]}", mustLimitJSON(t, delta)))
|
||||
_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(strings.Repeat(event, 6)), Chat, Responses, "", 512, 8, 512)
|
||||
if err == nil || !strings.Contains(err.Error(), "retained content") {
|
||||
t.Fatalf("%s: %v", field, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
func TestStreamBlockLimits(t *testing.T) {
|
||||
for _, kind := range []string{"text", "thinking", "reasoning", "summary", "tool"} {
|
||||
from := Anthropic
|
||||
in := limitEvent("{\"type\":\"message_start\",\"message\":{\"content\":[]}}")
|
||||
if kind == "reasoning" || kind == "summary" || kind == "tool" {
|
||||
from = Responses
|
||||
in = ""
|
||||
}
|
||||
if kind == "summary" {
|
||||
in = limitEvent("{\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"reasoning\",\"id\":\"r\"}}")
|
||||
}
|
||||
for i := 0; i < 4; i++ {
|
||||
m := object{"type": "content_block_start", "index": i, "content_block": object{"type": kind, "text": ""}}
|
||||
switch kind {
|
||||
case "thinking":
|
||||
m["content_block"] = object{"type": "thinking", "thinking": "", "signature": ""}
|
||||
case "reasoning", "tool":
|
||||
item := object{"type": "reasoning", "id": "r"}
|
||||
if kind == "tool" {
|
||||
item = object{"type": "function_call", "call_id": "c", "name": "f", "arguments": ""}
|
||||
}
|
||||
m = object{"type": "response.output_item.added", "output_index": i, "item": item}
|
||||
case "summary":
|
||||
m = object{"type": "response.reasoning_summary_part.added", "output_index": 0, "item_id": "r", "summary_index": i, "part": object{"type": "summary_text", "text": ""}}
|
||||
}
|
||||
in += limitEvent(mustLimitJSON(t, m))
|
||||
}
|
||||
_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in), from, Chat, "", 1024, 3, 1024)
|
||||
if err == nil || !strings.Contains(err.Error(), "3 blocks") {
|
||||
t.Fatalf("%s: %v", kind, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReasoningIdentityCumulativeLimit(t *testing.T) {
|
||||
var in strings.Builder
|
||||
for i := 0; i < 6; i++ {
|
||||
in.WriteString(limitEvent(mustLimitJSON(t, object{"type": "response.output_item.added", "output_index": i, "item": object{"type": "reasoning", "id": strings.Repeat("r", 100)}})))
|
||||
}
|
||||
_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in.String()), Responses, Chat, "", 512, 8, 512)
|
||||
if err == nil || !strings.Contains(err.Error(), "retained content") {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
func TestStreamBufferSemantics(t *testing.T) {
|
||||
var out bytes.Buffer
|
||||
s := &streamState{ctx: context.Background(), w: &out, from: Chat, to: Responses, source: map[int]*streamBlock{}, byteLimit: 64, blockLimit: 4}
|
||||
b, err := s.add(0, part{typ: "call"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.delta(b, "{\"a\":"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.Len() != 0 {
|
||||
t.Fatal("incomplete identity emitted")
|
||||
}
|
||||
b.id, b.name = "c", "f"
|
||||
if err = s.delta(b, "1}"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b.args != "{\"a\":1}" || b.sent != 7 {
|
||||
t.Fatal(b.part)
|
||||
}
|
||||
if err = s.close(b); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
text, err := s.add(1, part{typ: "text"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, v := range []string{"ab", "cd"} {
|
||||
if err = s.delta(text, v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if text.text != "abcd" {
|
||||
t.Fatal(text.text)
|
||||
}
|
||||
if err = s.finish("stop"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.end(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(out.String(), "abcd") || !strings.Contains(out.String(), "response.completed") {
|
||||
t.Fatal(out.String())
|
||||
}
|
||||
s.finished = false
|
||||
text.closed = false
|
||||
s.byteLimit = s.retained
|
||||
if err = s.delta(text, "x"); err == nil || text.text != "abcd" {
|
||||
t.Fatal("overflow changed content")
|
||||
}
|
||||
}
|
||||
func mustLimitJSON(t *testing.T, m object) string {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMergeUsageSizeBoundary(t *testing.T) {
|
||||
for _, kind := range []Kind{Chat, Responses, Anthropic} {
|
||||
t.Run(string(kind), func(t *testing.T) {
|
||||
m := object{"padding": strings.Repeat("x", maxStreamUsageBytes-len(`{"padding":""}`))}
|
||||
u, err := mergeUsage(kind, Usage{}, object{"usage": m})
|
||||
if err != nil || len(u.Raw) != maxStreamUsageBytes {
|
||||
t.Fatalf("exact limit: size=%d err=%v", len(u.Raw), err)
|
||||
}
|
||||
for _, update := range []object{
|
||||
{"usage": object{"padding": str(m["padding"]) + "x"}},
|
||||
{"usage": object{"new_key": 1}},
|
||||
} {
|
||||
got, err := mergeUsage(kind, u, update)
|
||||
if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") || !reflect.DeepEqual(got, u) {
|
||||
t.Fatalf("overflow: retained=%d err=%v", len(got.Raw), err)
|
||||
}
|
||||
}
|
||||
for i := 0; i < 3; i++ {
|
||||
got, err := mergeUsage(kind, u, object{"usage": m})
|
||||
if err != nil || !reflect.DeepEqual(got, u) {
|
||||
t.Fatalf("replacement at limit: %v", err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamUsageSizeLimits(t *testing.T) {
|
||||
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||
for _, cumulative := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("%s/%s/cumulative=%t", from, to, cumulative), func(t *testing.T) {
|
||||
var input strings.Builder
|
||||
writeUsage := func(u object) {
|
||||
m := object{"usage": u}
|
||||
switch from {
|
||||
case Anthropic:
|
||||
m["type"] = "ping"
|
||||
case Responses:
|
||||
m = object{"type": "response.in_progress", "response": m}
|
||||
}
|
||||
input.WriteString(limitEvent(mustLimitJSON(t, m)))
|
||||
}
|
||||
writeUsage(object{"prompt_tokens": 7, "input_tokens": 7})
|
||||
if cumulative {
|
||||
for i := 0; i < 32; i++ {
|
||||
writeUsage(object{fmt.Sprintf("key_%d", i): strings.Repeat("x", 1024)})
|
||||
}
|
||||
} else {
|
||||
writeUsage(object{"padding": strings.Repeat("x", maxStreamUsageBytes)})
|
||||
}
|
||||
u, err := Stream(context.Background(), io.Discard, strings.NewReader(input.String()), from, to, "")
|
||||
if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") {
|
||||
t.Fatalf("expected usage limit error, got %v", err)
|
||||
}
|
||||
if !u.Known || u.Input != 7 || len(u.Raw) > maxStreamUsageBytes {
|
||||
t.Fatalf("lost valid usage: known=%t input=%d size=%d", u.Known, u.Input, len(u.Raw))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeUsagePartialUpdates(t *testing.T) {
|
||||
u, err := mergeUsage(Anthropic, Usage{}, object{"message": object{"usage": object{"input_tokens": 5, "cache_read_input_tokens": 2}}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, err = mergeUsage(Anthropic, u, object{"usage": object{"output_tokens": 3}})
|
||||
if err != nil || !u.Known || u.Input != 7 || u.Output != 3 || u.CacheRead != 2 {
|
||||
t.Fatalf("partial update: %+v err=%v", u, err)
|
||||
}
|
||||
got, err := mergeUsage(Anthropic, u, object{})
|
||||
if err != nil || !reflect.DeepEqual(got, u) {
|
||||
t.Fatalf("missing usage changed totals: %+v err=%v", got, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package protocol
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMalformedInitialToolArguments(t *testing.T) {
|
||||
for _, bad := range []string{`null`, `[]`, `12`, `"not-an-object"`} {
|
||||
src := `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"tool1","name":"fn","input":` + bad + `}}` + "\n\n"
|
||||
_, err := Stream(context.Background(), io.Discard, strings.NewReader(src), Anthropic, Chat, "m")
|
||||
if err == nil || !strings.Contains(err.Error(), "tool input must be an object") {
|
||||
t.Fatalf("input %s: %v", bad, err)
|
||||
}
|
||||
}
|
||||
for _, field := range []string{``, `,"arguments":null`, `,"arguments":{}`, `,"arguments":12`} {
|
||||
src := `data: {"type":"response.output_item.added","output_index":0,"item":{"type":"function_call","call_id":"call1","name":"fn"` + field + `}}` + "\n\n"
|
||||
_, err := Stream(context.Background(), io.Discard, strings.NewReader(src), Responses, Chat, "m")
|
||||
if err == nil || !strings.Contains(err.Error(), "function arguments must be a string") {
|
||||
t.Fatalf("field %s: %v", field, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeNodeLimit(t *testing.T) {
|
||||
dense := `{"model":"m","pad":[` + strings.Repeat("0,", maxJSONNodes) + `0]}`
|
||||
if _, err := ConvertRequest(Chat, Chat, []byte(dense), "m"); err == nil || !strings.Contains(err.Error(), "exceeds") {
|
||||
t.Fatalf("dense body accepted: %v", err)
|
||||
}
|
||||
under := `{"model":"m","pad":[` + strings.Repeat("0,", maxJSONNodes/2) + `0]}`
|
||||
if _, err := ConvertRequest(Chat, Chat, []byte(under), "m"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := ConvertRequest(Chat, Chat, []byte(`{"broken":`), "m"); err == nil || strings.Contains(err.Error(), "exceeds") {
|
||||
t.Fatalf("malformed JSON misreported: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,265 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type channel struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"base_url"`
|
||||
Protocol string `json:"protocol"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Models map[string]string `json:"models"`
|
||||
APIKey string `json:"api_key,omitempty"`
|
||||
}
|
||||
type apiKey struct {
|
||||
Revoked bool `json:"revoked"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Prefix string `json:"prefix"`
|
||||
Enabled bool `json:"enabled"`
|
||||
AllChannels bool `json:"all_channels"`
|
||||
Channels []string `json:"channels"`
|
||||
ExpiresAt string `json:"expires_at"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
}
|
||||
|
||||
func (s *Server) loadChannels() ([]channel, error) {
|
||||
rows, e := s.db.Query("SELECT id,name,base_url,protocol,enabled,models,api_key FROM channels ORDER BY name")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []channel{}
|
||||
for rows.Next() {
|
||||
var c channel
|
||||
var m string
|
||||
if e = rows.Scan(&c.ID, &c.Name, &c.BaseURL, &c.Protocol, &c.Enabled, &m, &c.APIKey); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if e = json.Unmarshal([]byte(m), &c.Models); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
func (s *Server) channels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == "GET" {
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
for i := range cs {
|
||||
cs[i].APIKey = ""
|
||||
}
|
||||
writeJSON(w, cs)
|
||||
return
|
||||
}
|
||||
if r.Method != "POST" {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var c channel
|
||||
if decode(w, r, &c) != nil {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
u, e := url.Parse(c.BaseURL)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(c.BaseURL, "#") || c.Name == "" || strings.Contains(c.Name, "/") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
if c.Models == nil {
|
||||
c.Models = map[string]string{}
|
||||
}
|
||||
for a, m := range c.Models {
|
||||
if a == "" || m == "" {
|
||||
errJSON(w, 400, "invalid model mapping")
|
||||
return
|
||||
}
|
||||
}
|
||||
c.BaseURL = strings.TrimRight(c.BaseURL, "/")
|
||||
m, _ := json.Marshal(c.Models)
|
||||
if c.ID == "" {
|
||||
c.ID = randomID()
|
||||
_, e = s.db.Exec("INSERT INTO channels VALUES(?,?,?,?,?,?,?,?)", c.ID, c.Name, c.BaseURL, c.Protocol, c.Enabled, false, string(m), c.APIKey)
|
||||
} else {
|
||||
var old string
|
||||
if e = s.db.QueryRow("SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old); e == sql.ErrNoRows {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
} else if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if c.APIKey == "" {
|
||||
c.APIKey = old
|
||||
}
|
||||
_, e = s.db.Exec("UPDATE channels SET name=?,base_url=?,protocol=?,enabled=?,models=?,api_key=? WHERE id=?", c.Name, c.BaseURL, c.Protocol, c.Enabled, string(m), c.APIKey, c.ID)
|
||||
}
|
||||
if e != nil {
|
||||
errJSON(w, 400, "channel could not be saved; names must be unique")
|
||||
return
|
||||
}
|
||||
c.APIKey = ""
|
||||
writeJSON(w, c)
|
||||
}
|
||||
func (s *Server) channelID(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != "DELETE" {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
res, e := s.db.Exec("DELETE FROM channels WHERE id=?", strings.TrimPrefix(r.URL.Path, "/admin/api/channels/"))
|
||||
mutation(w, res, e)
|
||||
}
|
||||
func mutation(w http.ResponseWriter, res sql.Result, e error) {
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
n, e := res.RowsAffected()
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if n == 0 {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]bool{"ok": true})
|
||||
}
|
||||
func (s *Server) keys(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == "GET" {
|
||||
rows, e := s.db.Query("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at,revoked FROM keys ORDER BY created_at DESC")
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []apiKey{}
|
||||
for rows.Next() {
|
||||
var k apiKey
|
||||
var ch string
|
||||
if rows.Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt, &k.Revoked) != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if json.Unmarshal([]byte(ch), &k.Channels) != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
out = append(out, k)
|
||||
}
|
||||
if rows.Err() != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, out)
|
||||
return
|
||||
}
|
||||
if r.Method != "POST" {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var k apiKey
|
||||
if decode(w, r, &k) != nil || k.Name == "" {
|
||||
errJSON(w, 400, "invalid key")
|
||||
return
|
||||
}
|
||||
if k.ExpiresAt != "" {
|
||||
t, e := time.Parse(time.RFC3339, k.ExpiresAt)
|
||||
if e != nil || !t.After(time.Now()) {
|
||||
errJSON(w, 400, "invalid expiry")
|
||||
return
|
||||
}
|
||||
k.ExpiresAt = t.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
if k.Channels == nil {
|
||||
k.Channels = []string{}
|
||||
}
|
||||
for _, id := range k.Channels {
|
||||
var n int
|
||||
if s.db.QueryRow("SELECT count(*) FROM channels WHERE id=?", id).Scan(&n) != nil || n == 0 {
|
||||
errJSON(w, 400, "invalid channel permission")
|
||||
return
|
||||
}
|
||||
}
|
||||
raw := "sk-" + randomID()
|
||||
ch, _ := json.Marshal(k.Channels)
|
||||
_, e := s.db.Exec("INSERT INTO keys(id,name,prefix,hash,enabled,all_channels,channels,expires_at,created_at) VALUES(?,?,?,?,?,?,?,?,?)", randomID(), k.Name, raw[:11], hash(raw), true, k.AllChannels, string(ch), k.ExpiresAt, time.Now().UTC().Format(time.RFC3339Nano))
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]string{"key": raw})
|
||||
}
|
||||
func (s *Server) keyID(w http.ResponseWriter, r *http.Request) {
|
||||
id := strings.TrimPrefix(r.URL.Path, "/admin/api/keys/")
|
||||
if r.Method == "DELETE" {
|
||||
res, e := s.db.Exec("UPDATE keys SET enabled=0,revoked=1 WHERE id=?", id)
|
||||
mutation(w, res, e)
|
||||
return
|
||||
}
|
||||
if r.Method == "PATCH" {
|
||||
var x struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
if decode(w, r, &x) != nil || x.Enabled == nil {
|
||||
errJSON(w, 400, "invalid request")
|
||||
return
|
||||
}
|
||||
res, e := s.db.Exec("UPDATE keys SET enabled=? WHERE id=? AND revoked=0", *x.Enabled, id)
|
||||
mutation(w, res, e)
|
||||
return
|
||||
}
|
||||
errJSON(w, 405, "method not allowed")
|
||||
}
|
||||
func (s *Server) authenticate(r *http.Request) (apiKey, error) {
|
||||
var k apiKey
|
||||
raw := r.Header.Get("x-api-key")
|
||||
if a := r.Header.Get("Authorization"); a != "" {
|
||||
parts := strings.Fields(a)
|
||||
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
|
||||
return k, sql.ErrNoRows
|
||||
}
|
||||
if raw != "" && raw != parts[1] {
|
||||
return k, sql.ErrNoRows
|
||||
}
|
||||
raw = parts[1]
|
||||
}
|
||||
if raw == "" {
|
||||
return k, sql.ErrNoRows
|
||||
}
|
||||
var ch string
|
||||
e := s.db.QueryRow("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at FROM keys WHERE hash=? AND enabled=1 AND revoked=0", hash(raw)).Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt)
|
||||
if e != nil {
|
||||
return k, e
|
||||
}
|
||||
if k.ExpiresAt != "" {
|
||||
t, e := time.Parse(time.RFC3339Nano, k.ExpiresAt)
|
||||
if e != nil || !time.Now().Before(t) {
|
||||
return apiKey{}, sql.ErrNoRows
|
||||
}
|
||||
}
|
||||
e = json.Unmarshal([]byte(ch), &k.Channels)
|
||||
return k, e
|
||||
}
|
||||
func (k apiKey) allows(id string) bool {
|
||||
if k.AllChannels {
|
||||
return true
|
||||
}
|
||||
for _, v := range k.Channels {
|
||||
if v == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
discoveryTimeout = 30 * time.Second
|
||||
discoveryMaxBody = 1 << 20
|
||||
discoveryMaxPages = 20
|
||||
)
|
||||
|
||||
func (s *Server) discoverModels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var c channel
|
||||
if decode(w, r, &c) != nil {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
u, e := url.Parse(c.BaseURL)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || strings.Contains(c.BaseURL, "#") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
if c.ID != "" {
|
||||
var old string
|
||||
e = s.db.QueryRowContext(r.Context(), "SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old)
|
||||
if e == sql.ErrNoRows {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if c.APIKey == "" {
|
||||
c.APIKey = old
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), discoveryTimeout)
|
||||
defer cancel()
|
||||
models, e := s.fetchModels(ctx, c)
|
||||
if e != nil {
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||
errJSON(w, 504, "model discovery timed out")
|
||||
} else {
|
||||
errJSON(w, 502, "model discovery failed")
|
||||
}
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string][]string{"models": models})
|
||||
}
|
||||
|
||||
func (s *Server) fetchModels(ctx context.Context, c channel) ([]string, error) {
|
||||
// Share the bounded transport, but never forward credentials across redirects.
|
||||
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||
endpoint := strings.TrimRight(c.BaseURL, "/") + "/models"
|
||||
if c.Protocol == "anthropic" {
|
||||
endpoint = strings.TrimRight(c.BaseURL, "/") + "/v1/models"
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
cursors := map[string]bool{}
|
||||
after := ""
|
||||
for page := 0; page < discoveryMaxPages; page++ {
|
||||
u := endpoint
|
||||
if after != "" {
|
||||
u += "?" + url.Values{"after_id": {after}}.Encode()
|
||||
}
|
||||
req, e := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
if c.Protocol == "anthropic" {
|
||||
req.Header.Set("x-api-key", c.APIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
res, e := client.Do(req)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if res.StatusCode != http.StatusOK {
|
||||
res.Body.Close()
|
||||
return nil, errors.New("upstream status")
|
||||
}
|
||||
body, e := io.ReadAll(io.LimitReader(res.Body, discoveryMaxBody+1))
|
||||
res.Body.Close()
|
||||
if e != nil || len(body) > discoveryMaxBody {
|
||||
return nil, errors.New("upstream body")
|
||||
}
|
||||
var result struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
HasMore bool `json:"has_more"`
|
||||
LastID string `json:"last_id"`
|
||||
}
|
||||
if json.Unmarshal(body, &result) != nil || result.Data == nil {
|
||||
return nil, errors.New("upstream models")
|
||||
}
|
||||
for _, model := range result.Data {
|
||||
if strings.TrimSpace(model.ID) == "" {
|
||||
return nil, errors.New("upstream model id")
|
||||
}
|
||||
seen[model.ID] = true
|
||||
}
|
||||
if c.Protocol != "anthropic" || !result.HasMore {
|
||||
models := make([]string, 0, len(seen))
|
||||
for id := range seen {
|
||||
models = append(models, id)
|
||||
}
|
||||
sort.Strings(models)
|
||||
return models, nil
|
||||
}
|
||||
if result.LastID == "" || cursors[result.LastID] {
|
||||
return nil, errors.New("upstream pagination")
|
||||
}
|
||||
cursors[result.LastID] = true
|
||||
after = result.LastID
|
||||
}
|
||||
return nil, errors.New("upstream page limit")
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const discoveryPath = "/admin/api/channels/discover-models"
|
||||
|
||||
func discover(t *testing.T, s *Server, cookie *http.Cookie, c channel) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
b, e := json.Marshal(c)
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
return call(s, "POST", discoveryPath, string(b), cookie)
|
||||
}
|
||||
|
||||
func TestDiscoveryProtocolsAndReadOnly(t *testing.T) {
|
||||
for _, proto := range []string{"openai-chat", "openai-responses", "anthropic"} {
|
||||
t.Run(proto, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
calls := 0
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
if r.Method != "GET" || r.URL.Path != "/v1/models" {
|
||||
t.Error(r.Method, r.URL)
|
||||
}
|
||||
if proto == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "secret" || r.Header.Get("anthropic-version") != "2023-06-01" || r.Header.Get("Authorization") != "" {
|
||||
t.Error("wrong anthropic headers")
|
||||
}
|
||||
if r.URL.Query().Get("after_id") == "" {
|
||||
fmt.Fprint(w, `{"data":[{"id":"z"}],"has_more":true,"last_id":"cursor/&?"}`)
|
||||
return
|
||||
}
|
||||
if r.URL.Query().Get("after_id") != "cursor/&?" {
|
||||
t.Error(r.URL)
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer secret" || r.Header.Get("x-api-key") != "" {
|
||||
t.Error("wrong openai headers")
|
||||
}
|
||||
fmt.Fprint(w, `{"data":[{"id":"z"},{"id":"a"},{"id":"a"}]}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
draft := channel{BaseURL: up.URL + "/v1/", Protocol: proto, APIKey: "secret"}
|
||||
if proto == "anthropic" {
|
||||
draft.BaseURL = up.URL + "/"
|
||||
}
|
||||
w := discover(t, s, cookie, draft)
|
||||
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":["a","z"]}` {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
wantCalls := 1
|
||||
if proto == "anthropic" {
|
||||
wantCalls = 2
|
||||
}
|
||||
if calls != wantCalls {
|
||||
t.Fatal(calls)
|
||||
}
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil || len(cs) != 0 {
|
||||
t.Fatal(cs, e)
|
||||
}
|
||||
draft.Name = "saved"
|
||||
draft.Models = map[string]string{"alias": "a"}
|
||||
b, _ := json.Marshal(draft)
|
||||
w = call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var saved channel
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &saved); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
before, _ := s.loadChannels()
|
||||
draft.ID, draft.APIKey, draft.Name, draft.Models = saved.ID, "", "unsaved", nil
|
||||
if w = discover(t, s, cookie, draft); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
after, _ := s.loadChannels()
|
||||
if !reflect.DeepEqual(before, after) {
|
||||
t.Fatal("discovery changed channel")
|
||||
}
|
||||
replacement := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if proto == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "replacement" {
|
||||
t.Error("draft key not used")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer replacement" {
|
||||
t.Error("draft key not used")
|
||||
}
|
||||
fmt.Fprint(w, `{"data":[]}`)
|
||||
}))
|
||||
defer replacement.Close()
|
||||
draft.BaseURL, draft.APIKey = replacement.URL, "replacement"
|
||||
w = discover(t, s, cookie, draft)
|
||||
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":[]}` {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
after, _ = s.loadChannels()
|
||||
if !reflect.DeepEqual(before, after) {
|
||||
t.Fatal("draft replacement was saved")
|
||||
}
|
||||
var logs int
|
||||
if e := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&logs); e != nil || logs != 0 {
|
||||
t.Fatal(logs, e)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryAuthAndValidation(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
if w := call(s, "POST", discoveryPath, "{}", nil); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader("{}"))
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
r.Header.Del(header)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(header, w.Code)
|
||||
}
|
||||
}
|
||||
for _, method := range []string{"GET", "DELETE", "PUT"} {
|
||||
if w := call(s, method, discoveryPath, "{}", cookie); w.Code != 405 {
|
||||
t.Fatal(method, w.Code)
|
||||
}
|
||||
}
|
||||
for _, base := range []string{"", "ftp://example.com", "http://", "http://user:[email protected]", "http://example.com?x=1", "http://example.com?", "http://example.com#"} {
|
||||
if w := discover(t, s, cookie, channel{BaseURL: base, Protocol: "openai-chat"}); w.Code != 400 {
|
||||
t.Fatal(base, w.Code)
|
||||
}
|
||||
}
|
||||
if w := discover(t, s, cookie, channel{BaseURL: "http://example.com", Protocol: "unknown"}); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := discover(t, s, cookie, channel{ID: "missing", BaseURL: "http://example.com", Protocol: "anthropic"}); w.Code != 404 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, body := range []string{"{", "{} {}", strings.Repeat(" ", maxBody+1)} {
|
||||
if w := call(s, "POST", discoveryPath, body, cookie); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryUpstreamErrors(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status, calls int
|
||||
}{
|
||||
{"status", "secret upstream details", 401, 1},
|
||||
{"json", "secret", 200, 1},
|
||||
{"missing data", "{}", 200, 1},
|
||||
{"null", `{"data":null}`, 200, 1},
|
||||
{"empty id", `{"data":[{"id":""}]}`, 200, 1},
|
||||
{"trailing", `{"data":[]}{}`, 200, 1},
|
||||
{"oversized", strings.Repeat("x", discoveryMaxBody+1), 200, 1},
|
||||
{"missing cursor", `{"data":[],"has_more":true}`, 200, 1},
|
||||
{"repeated cursor", `{"data":[],"has_more":true,"last_id":"same"}`, 200, 2},
|
||||
{"page limit", "", 200, discoveryMaxPages},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
calls := 0
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
w.WriteHeader(tc.status)
|
||||
if tc.name == "page limit" {
|
||||
fmt.Fprintf(w, `{"data":[],"has_more":true,"last_id":"%d"}`, calls)
|
||||
} else {
|
||||
fmt.Fprint(w, tc.body)
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: "anthropic", APIKey: "secret"})
|
||||
if w.Code != 502 || strings.TrimSpace(w.Body.String()) != `{"error":"model discovery failed"}` || calls != tc.calls {
|
||||
t.Fatal(w.Code, w.Body.String(), calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryRedirectAndDeadline(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("redirect followed") }))
|
||||
defer target.Close()
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
||||
defer up.Close()
|
||||
for _, proto := range []string{"openai-chat", "anthropic"} {
|
||||
if w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: proto, APIKey: "secret"}); w.Code != 502 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
|
||||
defer slow.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
b, _ := json.Marshal(channel{BaseURL: slow.URL, Protocol: "openai-chat"})
|
||||
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader(string(b))).WithContext(ctx)
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 504 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelsListSelectedAliasesAndPermissions(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var allowed string
|
||||
for _, name := range []string{"selected", "forbidden", "disabled"} {
|
||||
b, _ := json.Marshal(channel{Name: name, BaseURL: "http://example.com", Protocol: "openai-chat", Enabled: name != "disabled", Models: map[string]string{"z": "vendor/real", "alias": "vendor/real"}})
|
||||
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var c channel
|
||||
json.Unmarshal(w.Body.Bytes(), &c)
|
||||
if name != "forbidden" {
|
||||
if allowed != "" {
|
||||
allowed += ","
|
||||
}
|
||||
allowed += fmt.Sprintf("%q", c.ID)
|
||||
}
|
||||
}
|
||||
w := call(s, "POST", "/admin/api/keys", `{"name":"client","channels":[`+allowed+`]}`, cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var key map[string]string
|
||||
json.Unmarshal(w.Body.Bytes(), &key)
|
||||
r := httptest.NewRequest("GET", "/v1/models", nil)
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
var result struct{ Data []struct{ ID string } }
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &result); e != nil || w.Code != 200 {
|
||||
t.Fatal(w.Code, e)
|
||||
}
|
||||
if len(result.Data) != 2 || result.Data[0].ID != "selected/alias" || result.Data[1].ID != "selected/z" {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
for _, model := range []string{"forbidden/alias", "disabled/alias"} {
|
||||
r = httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[]}`, model)))
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(model, w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type integrationKey struct{ raw, id string }
|
||||
|
||||
type integrationStreamRecorder struct {
|
||||
*httptest.ResponseRecorder
|
||||
flushed chan struct{}
|
||||
}
|
||||
|
||||
func (w *integrationStreamRecorder) Flush() {
|
||||
w.ResponseRecorder.Flush()
|
||||
select {
|
||||
case w.flushed <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func integrationAdmin(t *testing.T, s *Server, cookie *http.Cookie, method, path string, body any, out any) {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := call(s, method, path, string(b), cookie)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("%s %s: %d %s", method, path, w.Code, w.Body.String())
|
||||
}
|
||||
if out != nil {
|
||||
if err := json.Unmarshal(w.Body.Bytes(), out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func integrationChannel(t *testing.T, s *Server, cookie *http.Cookie, name, kind, base string) channel {
|
||||
t.Helper()
|
||||
if kind != "anthropic" {
|
||||
base += "/v1"
|
||||
}
|
||||
var ch channel
|
||||
integrationAdmin(t, s, cookie, "POST", "/admin/api/channels", channel{
|
||||
Name: name, Protocol: kind, BaseURL: base, Enabled: true,
|
||||
APIKey: "upstream-secret", Models: map[string]string{"gpt-5": "gpt-4", "gpt-4": "abc/gpt-4"},
|
||||
}, &ch)
|
||||
if ch.ID == "" || ch.APIKey != "" {
|
||||
t.Fatalf("invalid channel response: %+v", ch)
|
||||
}
|
||||
return ch
|
||||
}
|
||||
|
||||
func integrationCreateKey(t *testing.T, s *Server, cookie *http.Cookie, name string, ch channel) integrationKey {
|
||||
t.Helper()
|
||||
var created struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
integrationAdmin(t, s, cookie, "POST", "/admin/api/keys", apiKey{Name: name, Channels: []string{ch.ID}}, &created)
|
||||
var keys []apiKey
|
||||
integrationAdmin(t, s, cookie, "GET", "/admin/api/keys", nil, &keys)
|
||||
for _, k := range keys {
|
||||
if k.Name == name && created.Key != "" {
|
||||
return integrationKey{created.Key, k.ID}
|
||||
}
|
||||
}
|
||||
t.Fatal("created key missing from management API")
|
||||
return integrationKey{}
|
||||
}
|
||||
|
||||
func integrationRequest(kind, model string, stream bool) string {
|
||||
if kind == "openai-responses" {
|
||||
return fmt.Sprintf(`{"model":%q,"input":"hello","max_output_tokens":32,"stream":%t}`, model, stream)
|
||||
}
|
||||
return fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}],"max_tokens":32,"stream":%t}`, model, stream)
|
||||
}
|
||||
|
||||
func integrationPath(kind string) string {
|
||||
return map[string]string{"openai-chat": "/v1/chat/completions", "openai-responses": "/v1/responses", "anthropic": "/v1/messages"}[kind]
|
||||
}
|
||||
|
||||
func integrationProxy(s *Server, key integrationKey, kind, model string, stream bool) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", integrationPath(kind), strings.NewReader(integrationRequest(kind, model, stream)))
|
||||
if kind == "anthropic" {
|
||||
r.Header.Set("x-api-key", key.raw)
|
||||
} else {
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
|
||||
func integrationResponse(kind string) string {
|
||||
switch kind {
|
||||
case "openai-chat":
|
||||
return `{"id":"chat-test","object":"chat.completion","model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}`
|
||||
case "openai-responses":
|
||||
return `{"id":"resp-test","object":"response","model":"gpt-4","status":"completed","output":[{"type":"message","id":"msg-test","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello","annotations":[]}]}],"usage":{"input_tokens":12,"output_tokens":3,"input_tokens_details":{"cached_tokens":4}}}`
|
||||
default:
|
||||
return `{"id":"msg-test","type":"message","role":"assistant","model":"gpt-4","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":6,"output_tokens":3,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}`
|
||||
}
|
||||
}
|
||||
|
||||
func integrationSSE(kind string) string {
|
||||
switch kind {
|
||||
case "openai-chat":
|
||||
return "data: " + `{"id":"chat-test","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}` + "\n\ndata: " + `{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}` + "\n\ndata: [DONE]\n\n"
|
||||
case "openai-responses":
|
||||
return "event: response.created\ndata: " + `{"type":"response.created","response":{"id":"resp-test","status":"in_progress"}}` + "\n\n" +
|
||||
"event: response.output_item.added\ndata: " + `{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}` + "\n\n" +
|
||||
"event: response.output_text.delta\ndata: " + `{"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n" +
|
||||
"event: response.completed\ndata: " + `{"type":"response.completed","response":` + integrationResponse(kind) + "}\n\n"
|
||||
default:
|
||||
return "event: message_start\ndata: " + `{"type":"message_start","message":{"id":"msg-test","role":"assistant","content":[],"usage":{"input_tokens":6,"output_tokens":0,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}}` + "\n\n" +
|
||||
"event: content_block_start\ndata: " + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" +
|
||||
"event: content_block_delta\ndata: " + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}` + "\n\n" +
|
||||
"event: content_block_stop\ndata: " + `{"type":"content_block_stop","index":0}` + "\n\n" +
|
||||
"event: message_delta\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n" +
|
||||
"event: message_stop\ndata: " + `{"type":"message_stop"}` + "\n\n"
|
||||
}
|
||||
}
|
||||
|
||||
func integrationLogs(t *testing.T, s *Server, c *http.Cookie, key integrationKey) []requestLog {
|
||||
t.Helper()
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var logs []requestLog
|
||||
integrationAdmin(t, s, c, "GET", "/admin/api/logs?key_id="+key.id, nil, &logs)
|
||||
for _, l := range logs {
|
||||
if l.KeyID != key.id {
|
||||
t.Fatalf("cross-key log leak: %+v", l)
|
||||
}
|
||||
}
|
||||
return logs
|
||||
}
|
||||
|
||||
func integrationCount(t *testing.T, name string, got *int64, want int64) {
|
||||
t.Helper()
|
||||
if got == nil || *got != want {
|
||||
t.Fatalf("%s: got %v, want %d", name, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntegrationProtocolMatrix(t *testing.T) {
|
||||
kinds := []string{"openai-chat", "openai-responses", "anthropic"}
|
||||
for _, upKind := range kinds {
|
||||
for _, downKind := range kinds {
|
||||
t.Run(downKind+"_to_"+upKind, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var calls atomic.Int64
|
||||
var expectedModel atomic.Value
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
if r.URL.Path != integrationPath(upKind) {
|
||||
t.Errorf("upstream path: %s", r.URL.Path)
|
||||
}
|
||||
if upKind == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "upstream-secret" || r.Header.Get("anthropic-version") == "" {
|
||||
t.Error("missing Anthropic credentials/version")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
||||
t.Error("wrong upstream credentials")
|
||||
}
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
b, err := io.ReadAll(r.Body)
|
||||
if err != nil || json.Unmarshal(b, &body) != nil {
|
||||
t.Error("invalid upstream request")
|
||||
w.WriteHeader(400)
|
||||
return
|
||||
}
|
||||
if body.Model != expectedModel.Load() || !strings.Contains(string(b), "hello") {
|
||||
t.Errorf("upstream request: %s; want model %s", b, expectedModel.Load())
|
||||
}
|
||||
if body.Stream {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
for _, frame := range strings.Split(integrationSSE(upKind), "\n\n") {
|
||||
if frame != "" {
|
||||
fmt.Fprint(w, frame+"\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, integrationResponse(upKind))
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
abc := integrationChannel(t, s, cookie, "abc", upKind, up.URL)
|
||||
other := integrationChannel(t, s, cookie, "other", upKind, up.URL)
|
||||
keys := []integrationKey{integrationCreateKey(t, s, cookie, "first", abc), integrationCreateKey(t, s, cookie, "second", other)}
|
||||
for i, key := range keys {
|
||||
name, forbidden := "abc", "other"
|
||||
if i == 1 {
|
||||
name, forbidden = forbidden, name
|
||||
}
|
||||
for _, stream := range []bool{false, true} {
|
||||
for _, alias := range []string{"gpt-5", "gpt-4"} {
|
||||
expectedModel.Store("gpt-4")
|
||||
if alias == "gpt-4" {
|
||||
expectedModel.Store("abc/gpt-4")
|
||||
}
|
||||
w := integrationProxy(s, key, downKind, name+"/"+alias, stream)
|
||||
if w.Code != 200 || !strings.Contains(w.Body.String(), "hello") {
|
||||
t.Fatalf("stream=%t model=%s: %d %s", stream, alias, w.Code, w.Body.String())
|
||||
}
|
||||
if stream {
|
||||
terminal := map[string]string{"openai-chat": "data: [DONE]", "openai-responses": "event: response.completed", "anthropic": "event: message_stop"}[downKind]
|
||||
if !w.Flushed || !strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") || !strings.Contains(w.Body.String(), terminal) {
|
||||
t.Fatalf("incomplete SSE: %s", w.Body.String())
|
||||
}
|
||||
} else {
|
||||
var body map[string]any
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
field := map[string]string{"openai-chat": "choices", "openai-responses": "output", "anthropic": "content"}[downKind]
|
||||
if body[field] == nil {
|
||||
t.Fatalf("wrong downstream shape: %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
before := calls.Load()
|
||||
if w := integrationProxy(s, key, downKind, forbidden+"/gpt-5", false); w.Code != 403 {
|
||||
t.Fatalf("permission: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
if calls.Load() != before {
|
||||
t.Fatal("forbidden request reached upstream")
|
||||
}
|
||||
}
|
||||
if calls.Load() != 8 {
|
||||
t.Fatalf("upstream calls: %d", calls.Load())
|
||||
}
|
||||
for _, key := range keys {
|
||||
logs := integrationLogs(t, s, cookie, key)
|
||||
if len(logs) != 5 {
|
||||
t.Fatalf("logs: %+v", logs)
|
||||
}
|
||||
for _, l := range logs {
|
||||
if l.Status == 403 {
|
||||
if l.UsageKnown || l.Input != nil {
|
||||
t.Fatal("denied request has usage")
|
||||
}
|
||||
continue
|
||||
}
|
||||
actual := "gpt-4"
|
||||
if strings.HasSuffix(l.RequestModel, "/gpt-4") {
|
||||
actual = "abc/gpt-4"
|
||||
}
|
||||
if l.Status != 200 || l.Error != "" || !l.UsageKnown || l.ActualModel != actual || l.Downstream != downKind || l.Upstream != upKind {
|
||||
t.Fatalf("wrong log: %+v", l)
|
||||
}
|
||||
integrationCount(t, "input", l.Input, 12)
|
||||
integrationCount(t, "output", l.Output, 3)
|
||||
integrationCount(t, "cache read", l.CacheRead, 4)
|
||||
write := int64(0)
|
||||
if upKind == "anthropic" {
|
||||
write = 2
|
||||
}
|
||||
integrationCount(t, "cache write", l.CacheWrite, write)
|
||||
}
|
||||
var stats struct {
|
||||
Summary summary `json:"summary"`
|
||||
ByKey []struct {
|
||||
KeyID string `json:"key_id"`
|
||||
} `json:"by_key"`
|
||||
}
|
||||
integrationAdmin(t, s, cookie, "GET", "/admin/api/stats?key_id="+key.id, nil, &stats)
|
||||
if stats.Summary.Requests != 5 || stats.Summary.Success != 4 || stats.Summary.Failed != 1 || stats.Summary.Unknown != 1 || len(stats.ByKey) != 1 || stats.ByKey[0].KeyID != key.id {
|
||||
t.Fatalf("wrong isolated stats: %+v", stats)
|
||||
}
|
||||
integrationCount(t, "total input", stats.Summary.Input, 48)
|
||||
integrationCount(t, "total output", stats.Summary.Output, 12)
|
||||
integrationCount(t, "total cache read", stats.Summary.CacheRead, 16)
|
||||
write := int64(0)
|
||||
if upKind == "anthropic" {
|
||||
write = 8
|
||||
}
|
||||
integrationCount(t, "total cache write", stats.Summary.CacheWrite, write)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntegrationStreamFailures(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status int
|
||||
}{
|
||||
{"upstream_http_error", `{"error":"private upstream detail"}`, 429},
|
||||
{"error_event", "event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"private upstream detail\"}}\n\n", 200},
|
||||
{"malformed_event", "data: {invalid}\n\n", 200},
|
||||
{"truncated_stream", "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n", 200},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(tc.status)
|
||||
fmt.Fprint(w, tc.body)
|
||||
}))
|
||||
defer up.Close()
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "errors", ch)
|
||||
w := integrationProxy(s, key, "openai-responses", "abc/gpt-5", true)
|
||||
if w.Code != tc.status {
|
||||
t.Fatalf("status: %d body: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if strings.Contains(w.Body.String(), "response.completed") || strings.Contains(w.Body.String(), "private upstream detail") {
|
||||
t.Fatalf("false success or leaked upstream detail: %s", w.Body.String())
|
||||
}
|
||||
logs := integrationLogs(t, s, cookie, key)
|
||||
if len(logs) != 1 {
|
||||
t.Fatalf("logs: %+v", logs)
|
||||
}
|
||||
l := logs[0]
|
||||
if l.Status < 400 || l.Error == "" || l.UsageKnown || l.Input != nil || l.Output != nil {
|
||||
t.Fatalf("failure not recorded correctly: %+v", l)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicHeaderForwarding(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var mu sync.Mutex
|
||||
got := map[string]string{}
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
got["version"], got["beta"] = r.Header.Get("anthropic-version"), r.Header.Get("anthropic-beta")
|
||||
mu.Unlock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, integrationResponse("anthropic"))
|
||||
}))
|
||||
defer up.Close()
|
||||
ch := integrationChannel(t, s, cookie, "abc", "anthropic", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "beta", ch)
|
||||
request := func(version, beta string) {
|
||||
r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(integrationRequest("anthropic", "abc/gpt-5", false)))
|
||||
r.Header.Set("x-api-key", key.raw)
|
||||
if version != "" {
|
||||
r.Header.Set("anthropic-version", version)
|
||||
}
|
||||
if beta != "" {
|
||||
r.Header.Set("anthropic-beta", beta)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("status %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
request("2023-06-01", "context-1m-2025-08-07,interleaved-thinking-2025-05-14")
|
||||
mu.Lock()
|
||||
if got["version"] != "2023-06-01" || got["beta"] != "context-1m-2025-08-07,interleaved-thinking-2025-05-14" {
|
||||
t.Fatalf("valid headers dropped: %v", got)
|
||||
}
|
||||
mu.Unlock()
|
||||
request("bad version; injected", `weird"value`)
|
||||
mu.Lock()
|
||||
if got["version"] != "2023-06-01" || got["beta"] != "" {
|
||||
t.Fatalf("unsafe headers forwarded: %v", got)
|
||||
}
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
func TestIntegrationCancellationClosesUpstream(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
started, stopped := make(chan struct{}), make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
io.Copy(io.Discard, r.Body)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
close(started)
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
close(stopped)
|
||||
case <-release:
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
defer close(release)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "cancel", ch)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
r := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(integrationRequest("openai-responses", "abc/gpt-5", true))).WithContext(ctx)
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
done := make(chan struct{})
|
||||
w := &integrationStreamRecorder{ResponseRecorder: httptest.NewRecorder(), flushed: make(chan struct{}, 1)}
|
||||
go func() { s.ServeHTTP(w, r); close(done) }()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("upstream did not start")
|
||||
}
|
||||
select {
|
||||
case <-w.flushed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("downstream did not receive streaming output")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("proxy did not stop on cancellation")
|
||||
}
|
||||
select {
|
||||
case <-stopped:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("upstream request was not cancelled")
|
||||
}
|
||||
logs := integrationLogs(t, s, cookie, key)
|
||||
if len(logs) != 1 || logs[0].Status < 400 || logs[0].Error == "" || logs[0].UsageKnown || logs[0].Input != nil {
|
||||
t.Fatalf("cancel log: %+v", logs)
|
||||
}
|
||||
if strings.Contains(w.Body.String(), "response.completed") {
|
||||
t.Fatal("cancelled response marked completed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
maxLogFieldBytes = 512
|
||||
maxLogErrorBytes = 4096
|
||||
maxLogUsageBytes = 16384
|
||||
logCleanupBatch = 1000
|
||||
)
|
||||
|
||||
func (s *Server) SetLogRetention(days int) {
|
||||
if days < 0 {
|
||||
days = 0
|
||||
}
|
||||
s.logRetentionDays.Store(int64(days))
|
||||
}
|
||||
|
||||
func truncateLogField(value string, limit int) string {
|
||||
if len(value) <= limit {
|
||||
return value
|
||||
}
|
||||
for limit > 0 && !utf8.RuneStart(value[limit]) {
|
||||
limit--
|
||||
}
|
||||
return strings.Clone(value[:limit])
|
||||
}
|
||||
|
||||
func boundedLog(l requestLog) requestLog {
|
||||
for _, field := range []*string{&l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream} {
|
||||
*field = truncateLogField(*field, maxLogFieldBytes)
|
||||
}
|
||||
l.Error = truncateLogField(l.Error, maxLogErrorBytes)
|
||||
if len(l.Usage.Raw) > maxLogUsageBytes {
|
||||
// Keep valid JSON without retaining an attacker-controlled oversized object.
|
||||
l.Usage.Raw = []byte(`{"truncated":true}`)
|
||||
}
|
||||
l.RawUsage = nil
|
||||
return l
|
||||
}
|
||||
|
||||
func (s *Server) purgeExpiredLogs(ctx context.Context, days int) error {
|
||||
if days <= 0 {
|
||||
return nil
|
||||
}
|
||||
cutoff := time.Now().UTC().AddDate(0, 0, -days).Truncate(time.Second)
|
||||
_, err := s.db.ExecContext(ctx, "DELETE FROM logs WHERE id IN (SELECT id FROM logs WHERE created_at<? AND julianday(created_at)<julianday(?) LIMIT ?)", cutoff.Format("2006-01-02T15:04:05"), cutoff.Format(time.RFC3339), logCleanupBatch)
|
||||
return err
|
||||
}
|
||||
|
||||
type logEntry struct {
|
||||
value requestLog
|
||||
barrier chan error
|
||||
}
|
||||
|
||||
func (s *Server) record(l requestLog) {
|
||||
s.logMu.RLock()
|
||||
defer s.logMu.RUnlock()
|
||||
if s.logsClosed {
|
||||
return
|
||||
}
|
||||
l = boundedLog(l)
|
||||
l.Usage.Raw = append([]byte(nil), l.Usage.Raw...)
|
||||
s.logQueue <- logEntry{value: l}
|
||||
}
|
||||
|
||||
func (s *Server) flushLogs() error {
|
||||
s.logMu.RLock()
|
||||
if s.logsClosed {
|
||||
s.logMu.RUnlock()
|
||||
<-s.logDone
|
||||
return s.logErr
|
||||
}
|
||||
done := make(chan error, 1)
|
||||
s.logQueue <- logEntry{barrier: done}
|
||||
s.logMu.RUnlock()
|
||||
return <-done
|
||||
}
|
||||
|
||||
func (s *Server) writeLogs() {
|
||||
defer close(s.logDone)
|
||||
ticker := time.NewTicker(100 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
cleanup := time.NewTicker(time.Minute)
|
||||
defer cleanup.Stop()
|
||||
batch := make([]requestLog, 0, 100)
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
for attempt := 0; ; attempt++ {
|
||||
if s.persistLogs(batch) == nil {
|
||||
break
|
||||
}
|
||||
if attempt == 3 {
|
||||
s.logErr = errors.New("request log persistence failed after 3 retries")
|
||||
log.Printf("ai-router: request log persistence failed; batch_size=%d retries=3", len(batch))
|
||||
break
|
||||
}
|
||||
log.Printf("ai-router: retrying request log persistence; batch_size=%d retry=%d", len(batch), attempt+1)
|
||||
time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
|
||||
}
|
||||
clear(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case entry, ok := <-s.logQueue:
|
||||
if !ok {
|
||||
flush()
|
||||
return
|
||||
}
|
||||
if entry.barrier != nil {
|
||||
flush()
|
||||
entry.barrier <- s.logErr
|
||||
continue
|
||||
}
|
||||
batch = append(batch, entry.value)
|
||||
if len(batch) == cap(batch) {
|
||||
flush()
|
||||
}
|
||||
case <-ticker.C:
|
||||
flush()
|
||||
case <-cleanup.C:
|
||||
if days := s.logRetentionDays.Load(); days > 0 {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
err := s.purgeExpiredLogs(ctx, int(days))
|
||||
cancel()
|
||||
if err != nil {
|
||||
log.Print("ai-router: request log retention cleanup failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) persistLogs(batch []requestLog) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
for _, l := range batch {
|
||||
if err := insertLog(ctx, tx, l); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func waitLogCondition(t *testing.T, condition func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for !condition() {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("condition timed out")
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func logCount(t *testing.T, db *sql.DB) int {
|
||||
t.Helper()
|
||||
var n int
|
||||
if err := db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&n); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func TestLogAsyncAndRawCopy(t *testing.T) {
|
||||
s := testServer(t)
|
||||
conn, err := s.db.Conn(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
raw := []byte(`{"input_tokens":12}`)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.record(requestLog{Usage: protocol.Usage{Raw: raw}})
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("record waited for database")
|
||||
}
|
||||
raw[0] = 'x'
|
||||
conn.Close()
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got string
|
||||
if err := s.db.QueryRow("SELECT raw_usage FROM logs").Scan(&got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != `{"input_tokens":12}` {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogTimerAndBatchDrain(t *testing.T) {
|
||||
s := testServer(t)
|
||||
s.record(requestLog{})
|
||||
waitLogCondition(t, func() bool { return logCount(t, s.db) == 1 })
|
||||
for i := 0; i < 257; i++ {
|
||||
s.record(requestLog{})
|
||||
}
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != 258 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogBackpressureAndConcurrentClose(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "logs.db")
|
||||
s, err := New(path, "password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { s.Close() })
|
||||
conn, err := s.db.Conn(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
s.record(requestLog{})
|
||||
waitLogCondition(t, func() bool { return s.db.Stats().WaitCount > 0 })
|
||||
if cap(s.logQueue) != 1024 {
|
||||
t.Fatal(cap(s.logQueue))
|
||||
}
|
||||
for i := 0; i < 1024; i++ {
|
||||
s.record(requestLog{})
|
||||
}
|
||||
enqueued := make(chan struct{})
|
||||
go func() { s.record(requestLog{}); close(enqueued) }()
|
||||
select {
|
||||
case <-enqueued:
|
||||
t.Fatal("full queue did not apply backpressure")
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if err := s.Close(); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
conn.Close()
|
||||
done := make(chan struct{})
|
||||
go func() { wg.Wait(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Close deadlocked")
|
||||
}
|
||||
<-enqueued
|
||||
s.record(requestLog{})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db, err := sql.Open("sqlite", path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
if n := logCount(t, db); n != 1026 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogPersistenceErrorIsStickyAndRedacted(t *testing.T) {
|
||||
s := testServer(t)
|
||||
var output bytes.Buffer
|
||||
previous := log.Writer()
|
||||
log.SetOutput(&output)
|
||||
defer log.SetOutput(previous)
|
||||
if _, err := s.db.Exec(`CREATE TRIGGER reject_log BEFORE INSERT ON logs WHEN NEW.key_name='secret' BEGIN SELECT RAISE(FAIL,'private database secret'); END`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.record(requestLog{KeyName: "ok"})
|
||||
s.record(requestLog{KeyName: "secret"})
|
||||
start := time.Now()
|
||||
err := s.flushLogs()
|
||||
if err == nil {
|
||||
t.Fatal("missing persistence error")
|
||||
}
|
||||
if time.Since(start) > time.Second {
|
||||
t.Fatal("unbounded retries")
|
||||
}
|
||||
if n := logCount(t, s.db); n != 0 {
|
||||
t.Fatalf("partial transaction: %d", n)
|
||||
}
|
||||
if strings.Contains(output.String(), "secret") || strings.Contains(err.Error(), "secret") {
|
||||
t.Fatal("error leaked metadata")
|
||||
}
|
||||
if n := strings.Count(output.String(), "retrying"); n != 3 {
|
||||
t.Fatalf("retries: %d", n)
|
||||
}
|
||||
if _, err := s.db.Exec("DROP TRIGGER reject_log"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.record(requestLog{})
|
||||
if err := s.flushLogs(); err == nil {
|
||||
t.Fatal("error was not sticky")
|
||||
}
|
||||
if n := logCount(t, s.db); n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
if err := s.Close(); err == nil {
|
||||
t.Fatal("Close lost persistence error")
|
||||
}
|
||||
if err := s.Close(); err == nil {
|
||||
t.Fatal("second Close lost error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseWaitsForActiveRequest(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "logs.db")
|
||||
s, err := New(path, "test-password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { s.Close() })
|
||||
started, release := make(chan struct{}), make(chan struct{})
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
close(started)
|
||||
<-release
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(integrationResponse("openai-chat")))
|
||||
}))
|
||||
defer up.Close()
|
||||
defer func() {
|
||||
select {
|
||||
case <-release:
|
||||
default:
|
||||
close(release)
|
||||
}
|
||||
}()
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "active", ch)
|
||||
requestDone := make(chan struct{})
|
||||
go func() { integrationProxy(s, key, "openai-chat", "abc/gpt-5", false); close(requestDone) }()
|
||||
<-started
|
||||
closed := make(chan error, 1)
|
||||
go func() { closed <- s.Close() }()
|
||||
waitLogCondition(t, func() bool { s.lifecycleMu.Lock(); defer s.lifecycleMu.Unlock(); return s.closing })
|
||||
if w := call(s, "GET", "/v1/models", "", nil); w.Code != 503 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
select {
|
||||
case <-closed:
|
||||
t.Fatal("Close skipped active request")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
<-requestDone
|
||||
select {
|
||||
case err := <-closed:
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not drain")
|
||||
}
|
||||
db, err := sql.Open("sqlite", path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
if n := logCount(t, db); n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const modelTestTimeout = 30 * time.Second
|
||||
const modelTestMaxBody = 1 << 20
|
||||
|
||||
type modelTestResult struct {
|
||||
OK bool `json:"ok"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
Status int `json:"status"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// testModel is an administrative diagnostic, excluded from call logs and usage
|
||||
// statistics. It creates no API key; upstream providers may still bill the probe.
|
||||
func (s *Server) testModel(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
Model string `json:"model"`
|
||||
Prompt *string `json:"prompt"`
|
||||
}
|
||||
if decode(w, r, &input) != nil {
|
||||
errJSON(w, 400, "invalid request")
|
||||
return
|
||||
}
|
||||
name, alias, ok := strings.Cut(input.Model, "/")
|
||||
prompt := "Reply with OK."
|
||||
if input.Prompt != nil {
|
||||
prompt = *input.Prompt
|
||||
}
|
||||
if !ok || name == "" || alias == "" || strings.TrimSpace(prompt) == "" || utf8.RuneCountInString(prompt) > 2000 {
|
||||
errJSON(w, 400, "invalid model or prompt")
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), modelTestTimeout)
|
||||
defer cancel()
|
||||
var c channel
|
||||
var models string
|
||||
e := s.db.QueryRowContext(ctx, "SELECT base_url,protocol,enabled,models,api_key FROM channels WHERE name=?", name).Scan(&c.BaseURL, &c.Protocol, &c.Enabled, &models, &c.APIKey)
|
||||
if e == sql.ErrNoRows {
|
||||
errJSON(w, 403, "model or channel unavailable")
|
||||
return
|
||||
}
|
||||
if e != nil || json.Unmarshal([]byte(models), &c.Models) != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
actual, mapped := c.Models[alias]
|
||||
if !c.Enabled || !mapped || actual == "" {
|
||||
errJSON(w, 403, "model or channel unavailable")
|
||||
return
|
||||
}
|
||||
writeJSON(w, s.probeModel(ctx, c, actual, prompt))
|
||||
}
|
||||
|
||||
func (s *Server) probeModel(ctx context.Context, c channel, actual, prompt string) (result modelTestResult) {
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
result.DurationMS = time.Since(start).Milliseconds()
|
||||
if ctx.Err() != nil {
|
||||
result.OK = false
|
||||
result.Error = "model test canceled"
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
result.Error = "model test timed out"
|
||||
}
|
||||
}
|
||||
}()
|
||||
result.Error = "upstream request failed"
|
||||
up := protocol.Kind(c.Protocol)
|
||||
body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"role": "user", "content": prompt}}, "stream": false, "max_completion_tokens": 128})
|
||||
payload, e := protocol.ConvertRequest(protocol.Chat, up, body, actual)
|
||||
if e != nil {
|
||||
result.Error = "invalid upstream configuration"
|
||||
return
|
||||
}
|
||||
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
||||
req, e := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+suffix, bytes.NewReader(payload))
|
||||
if e != nil {
|
||||
result.Error = "invalid upstream configuration"
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if up == protocol.Anthropic {
|
||||
req.Header.Set("x-api-key", c.APIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||
res, e := client.Do(req)
|
||||
if e != nil {
|
||||
if errors.Is(e, context.DeadlineExceeded) {
|
||||
result.Error = "model test timed out"
|
||||
}
|
||||
return
|
||||
}
|
||||
defer res.Body.Close()
|
||||
result.Status = res.StatusCode
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
result.Error = "upstream rejected request"
|
||||
return
|
||||
}
|
||||
body, e = io.ReadAll(io.LimitReader(res.Body, modelTestMaxBody+1))
|
||||
if len(body) > modelTestMaxBody {
|
||||
result.Error = "upstream response too large"
|
||||
return
|
||||
}
|
||||
if e != nil || !validModelTestResponse(up, body) {
|
||||
result.Error = "invalid upstream response"
|
||||
return
|
||||
}
|
||||
result.OK, result.Error = true, ""
|
||||
return
|
||||
}
|
||||
|
||||
func validModelTestResponse(kind protocol.Kind, body []byte) bool {
|
||||
var r struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Error json.RawMessage `json:"error"`
|
||||
Status string `json:"status"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
Incomplete struct {
|
||||
Reason string `json:"reason"`
|
||||
} `json:"incomplete_details"`
|
||||
Output []json.RawMessage `json:"output"`
|
||||
Content []json.RawMessage `json:"content"`
|
||||
Choices []struct {
|
||||
Finish string `json:"finish_reason"`
|
||||
Message struct {
|
||||
Role string `json:"role"`
|
||||
Content *string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if json.Unmarshal(body, &r) != nil || r.ID == "" || (len(r.Error) != 0 && string(r.Error) != "null") {
|
||||
return false
|
||||
}
|
||||
switch kind {
|
||||
case protocol.Chat:
|
||||
if r.Object != "chat.completion" || len(r.Choices) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, c := range r.Choices {
|
||||
if c.Message.Role != "assistant" || (c.Finish != "stop" && c.Finish != "length") || (c.Message.Content == nil && c.Finish != "length") {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case protocol.Responses:
|
||||
if r.Object != "response" || r.Output == nil {
|
||||
return false
|
||||
}
|
||||
exhausted := r.Status == "incomplete" && r.Incomplete.Reason == "max_output_tokens"
|
||||
if !exhausted && (r.Status != "completed" || len(r.Output) == 0) {
|
||||
return false
|
||||
}
|
||||
for _, raw := range r.Output {
|
||||
var item struct {
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []json.RawMessage `json:"content"`
|
||||
Summary []json.RawMessage `json:"summary"`
|
||||
}
|
||||
if json.Unmarshal(raw, &item) != nil {
|
||||
return false
|
||||
}
|
||||
if item.Type == "reasoning" && item.Summary != nil {
|
||||
continue
|
||||
}
|
||||
if item.Type != "message" || item.Role != "assistant" || !validProbeBlocks(item.Content, "output_text") {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case protocol.Anthropic:
|
||||
if r.Type != "message" || r.Role != "assistant" || r.Content == nil {
|
||||
return false
|
||||
}
|
||||
if r.StopReason == "max_tokens" && len(r.Content) == 0 {
|
||||
return true
|
||||
}
|
||||
return (r.StopReason == "end_turn" || r.StopReason == "max_tokens") && validProbeBlocks(r.Content, "text")
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func validProbeBlocks(blocks []json.RawMessage, typ string) bool {
|
||||
if len(blocks) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, raw := range blocks {
|
||||
var b struct {
|
||||
Type string `json:"type"`
|
||||
Text *string `json:"text"`
|
||||
Thinking *string `json:"thinking"`
|
||||
Signature *string `json:"signature"`
|
||||
Data *string `json:"data"`
|
||||
}
|
||||
if json.Unmarshal(raw, &b) != nil {
|
||||
return false
|
||||
}
|
||||
if typ == "text" && ((b.Type == "thinking" && b.Thinking != nil && b.Signature != nil) || (b.Type == "redacted_thinking" && b.Data != nil)) {
|
||||
continue
|
||||
}
|
||||
if b.Type != typ || b.Text == nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const probePath = "/admin/api/models/test"
|
||||
const probeChat = `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":"OK"},"finish_reason":"stop"}]}`
|
||||
|
||||
func saveProbeChannel(t *testing.T, s *Server, cookie *http.Cookie, url, kind string, enabled bool) {
|
||||
t.Helper()
|
||||
b, _ := json.Marshal(channel{Name: "saved", BaseURL: url, Protocol: kind, Enabled: enabled, Models: map[string]string{"alias": "actual/model"}, APIKey: "secret-key"})
|
||||
if w := call(s, "POST", "/admin/api/channels", string(b), cookie); w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func probeResult(t *testing.T, w *httptest.ResponseRecorder) modelTestResult {
|
||||
t.Helper()
|
||||
var result modelTestResult
|
||||
if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &result) != nil {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
if result.DurationMS < 0 || strings.Contains(w.Body.String(), "secret-key") {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func TestModelProbeProtocols(t *testing.T) {
|
||||
for _, tc := range []struct{ kind, path, limit, response string }{
|
||||
{"openai-chat", "/chat/completions", "max_completion_tokens", probeChat},
|
||||
{"openai-responses", "/responses", "max_output_tokens", `{"id":"r","object":"response","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}]}`},
|
||||
{"anthropic", "/v1/messages", "max_tokens", `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{"type":"text","text":"OK"}]}`},
|
||||
} {
|
||||
t.Run(tc.kind, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
expectedPrompt := "Reply with OK."
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
var obj map[string]any
|
||||
if json.Unmarshal(b, &obj) != nil || obj["model"] != "actual/model" || obj[tc.limit] != float64(128) || obj["stream"] != false || !strings.Contains(string(b), expectedPrompt) {
|
||||
t.Error(string(b))
|
||||
}
|
||||
if tc.kind == "openai-chat" && obj["max_tokens"] != nil {
|
||||
t.Error("legacy token limit")
|
||||
}
|
||||
if r.Method != "POST" || r.URL.Path != tc.path {
|
||||
t.Error(r.Method, r.URL.Path)
|
||||
}
|
||||
if tc.kind == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "secret-key" || r.Header.Get("anthropic-version") != "2023-06-01" {
|
||||
t.Error("headers")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer secret-key" {
|
||||
t.Error("authorization")
|
||||
}
|
||||
fmt.Fprint(w, tc.response)
|
||||
}))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, tc.kind, true)
|
||||
for _, prompt := range []string{"", strings.Repeat("中", 2000)} {
|
||||
input := map[string]string{"model": "saved/alias"}
|
||||
if prompt != "" {
|
||||
expectedPrompt = prompt
|
||||
input["prompt"] = prompt
|
||||
}
|
||||
b, _ := json.Marshal(input)
|
||||
result := probeResult(t, call(s, "POST", probePath, string(b), cookie))
|
||||
if !result.OK || result.Status != 200 || result.Error != "" {
|
||||
t.Fatal(result)
|
||||
}
|
||||
}
|
||||
for _, table := range []string{"keys", "logs"} {
|
||||
var n int
|
||||
if e := s.db.QueryRow("SELECT count(*) FROM " + table).Scan(&n); e != nil || n != 0 {
|
||||
t.Fatal(table, n, e)
|
||||
}
|
||||
}
|
||||
if len(s.logQueue) != 0 {
|
||||
t.Fatal("queued formal log")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelProbeGuardAndInput(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
for _, prompt := range []string{"", " ", "\t\r\n", "\u3000\u00a0"} {
|
||||
body, _ := json.Marshal(map[string]string{"model": "saved/alias", "prompt": prompt})
|
||||
if w := call(s, "POST", probePath, string(body), cookie); w.Code != 400 {
|
||||
t.Fatalf("blank prompt %q: status %d", prompt, w.Code)
|
||||
}
|
||||
}
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("unexpected upstream request") }))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", false)
|
||||
if w := call(s, "POST", probePath, `{"model":"saved/alias"}`, nil); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`))
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
r.Header.Del(header)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
for _, body := range []string{`{"model":"saved/alias"}`, `{"model":"missing/alias"}`} {
|
||||
if w := call(s, "POST", probePath, body, cookie); w.Code != 403 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
if _, e := s.db.Exec("UPDATE channels SET enabled=1"); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if w := call(s, "POST", probePath, `{"model":"saved/unmapped"}`, cookie); w.Code != 403 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, body := range []string{`{}`, `null`, `{"model":"alias"}`, `{"model":"/alias"}`, `{"model":"saved/"}`, `{"model":"saved/alias","prompt":2}`, `{"model":"saved/alias"} {}`, `{"model":"saved/alias","prompt":"` + strings.Repeat("中", 2001) + `"}`} {
|
||||
if w := call(s, "POST", probePath, body, cookie); w.Code != 400 {
|
||||
t.Fatal(body, w.Code)
|
||||
}
|
||||
}
|
||||
if w := call(s, "GET", probePath, "", cookie); w.Code != 405 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelProbeFailures(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status int
|
||||
error string
|
||||
}{
|
||||
{"rejected", "secret-key upstream body", 401, "upstream rejected request"},
|
||||
{"redirect", "secret-key", 302, "upstream rejected request"},
|
||||
{"html", "<html>secret-key</html>", 200, "invalid upstream response"},
|
||||
{"empty", "{}", 200, "invalid upstream response"},
|
||||
{"oversize", strings.Repeat("x", modelTestMaxBody+1), 200, "upstream response too large"},
|
||||
{"limit", probeChat + strings.Repeat(" ", modelTestMaxBody-len(probeChat)), 200, ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
t.Error("followed redirect")
|
||||
}
|
||||
w.Header().Set("Location", "/redirected")
|
||||
w.WriteHeader(tc.status)
|
||||
fmt.Fprint(w, tc.body)
|
||||
}))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
||||
result := probeResult(t, call(s, "POST", probePath, `{"model":"saved/alias"}`, cookie))
|
||||
if result.Status != tc.status || result.Error != tc.error || result.OK != (tc.error == "") {
|
||||
t.Fatal(result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelProbeResponseValidation(t *testing.T) {
|
||||
if !validModelTestResponse(protocol.Anthropic, []byte(`{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[{"type":"thinking","thinking":"reasoning","signature":"sig"}]}`)) {
|
||||
t.Fatal("valid thinking-only response rejected")
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
kind protocol.Kind
|
||||
body string
|
||||
ok bool
|
||||
}{
|
||||
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":null},"finish_reason":"length"}]}`, true},
|
||||
{protocol.Responses, `{"id":"r","object":"response","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"output":[{"type":"reasoning","summary":[]}]}`, true},
|
||||
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[]}`, true},
|
||||
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{}]}`, false},
|
||||
{protocol.Responses, `{"id":"r","object":"response","status":"failed","output":[]}`, false},
|
||||
{protocol.Responses, `{"id":"r","object":"response","status":"completed","output":[{}]}`, false},
|
||||
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{}]}`, false},
|
||||
} {
|
||||
if got := validModelTestResponse(tc.kind, []byte(tc.body)); got != tc.ok {
|
||||
t.Error(tc.kind, tc.body, got)
|
||||
}
|
||||
}
|
||||
for _, kind := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
||||
for _, body := range []string{"{}", "null", "[]", probeChat + "{}", `{"error":{"message":"secret-key"}}`} {
|
||||
if validModelTestResponse(kind, []byte(body)) {
|
||||
t.Error(kind, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type probeRoundTripper func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f probeRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
|
||||
func TestModelProbeTimeoutAndCancellation(t *testing.T) {
|
||||
for _, cancelRequest := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(cancelRequest), func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
s.client.Transport = probeRoundTripper(func(r *http.Request) (*http.Response, error) {
|
||||
deadline, ok := r.Context().Deadline()
|
||||
if !ok || time.Until(deadline) > 30*time.Second || time.Until(deadline) < 29*time.Second {
|
||||
t.Error("missing 30s deadline", deadline)
|
||||
}
|
||||
if cancelRequest {
|
||||
cancel()
|
||||
<-r.Context().Done()
|
||||
return nil, r.Context().Err()
|
||||
}
|
||||
return nil, context.DeadlineExceeded
|
||||
})
|
||||
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`)).WithContext(ctx)
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
result := probeResult(t, w)
|
||||
if result.OK || result.Status != 0 || result.Error == "" {
|
||||
t.Fatal(result)
|
||||
}
|
||||
})
|
||||
}
|
||||
s := testServer(t)
|
||||
canceled := make(chan struct{})
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(200)
|
||||
w.(http.Flusher).Flush()
|
||||
<-r.Context().Done()
|
||||
close(canceled)
|
||||
}))
|
||||
defer up.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
result := s.probeModel(ctx, channel{BaseURL: up.URL, Protocol: "openai-chat"}, "actual", "test")
|
||||
if result.OK || result.Status != 200 || result.Error != "model test timed out" {
|
||||
t.Fatal(result)
|
||||
}
|
||||
select {
|
||||
case <-canceled:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("upstream was not canceled")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Explicitly opt in: ROUTER_PERF=1 go test ./internal/server -run TestPerformanceHTTP -v -count=1.
|
||||
func TestPerformanceHTTP(t *testing.T) {
|
||||
if os.Getenv("ROUTER_PERF") != "1" {
|
||||
t.Skip("opt-in load test")
|
||||
}
|
||||
for _, stream := range []bool{false, true} {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
io.Copy(io.Discard, r.Body)
|
||||
r.Body.Close()
|
||||
if stream {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
for _, f := range strings.Split(integrationSSE("openai-chat"), "\n\n") {
|
||||
if f != "" {
|
||||
fmt.Fprint(w, f+"\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, integrationResponse("openai-chat"))
|
||||
}
|
||||
}))
|
||||
for _, mode := range []string{"direct", "memory", "disk"} {
|
||||
for _, down := range []string{"openai-chat", "anthropic"} {
|
||||
if mode == "direct" && down != "openai-chat" {
|
||||
continue
|
||||
}
|
||||
for _, concurrency := range []int{1, 8, 32} {
|
||||
t.Run(fmt.Sprintf("stream=%t/%s/%s/c%d", stream, mode, down, concurrency), func(t *testing.T) {
|
||||
target := up.URL + integrationPath("openai-chat")
|
||||
raw := "test-key"
|
||||
var s *Server
|
||||
var router *httptest.Server
|
||||
if mode != "direct" {
|
||||
path := ":memory:"
|
||||
if mode == "disk" {
|
||||
path = filepath.Join(t.TempDir(), "router.db")
|
||||
}
|
||||
var err error
|
||||
s, err = New(path, "test-password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetLimits(1000000, 64, 128); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "perf", ch)
|
||||
raw = key.raw
|
||||
router = httptest.NewServer(s)
|
||||
defer router.Close()
|
||||
target = router.URL + integrationPath(down)
|
||||
}
|
||||
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||
tr.MaxIdleConnsPerHost = 64
|
||||
defer tr.CloseIdleConnections()
|
||||
client := &http.Client{Transport: tr, Timeout: 10 * time.Second}
|
||||
body := integrationRequest(down, "abc/gpt-5", stream)
|
||||
request := func() (time.Duration, time.Duration, bool) {
|
||||
start := time.Now()
|
||||
r, _ := http.NewRequest("POST", target, strings.NewReader(body))
|
||||
r.Header.Set("Authorization", "Bearer "+raw)
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(r)
|
||||
if err != nil {
|
||||
return time.Since(start), 0, false
|
||||
}
|
||||
first := make([]byte, 1)
|
||||
_, err = resp.Body.Read(first)
|
||||
ttfb := time.Since(start)
|
||||
_, e := io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
return time.Since(start), ttfb, resp.StatusCode == 200 && (err == nil || err == io.EOF) && e == nil
|
||||
}
|
||||
for i := 0; i < 20; i++ {
|
||||
if _, _, ok := request(); !ok {
|
||||
t.Fatal("warmup failed")
|
||||
}
|
||||
}
|
||||
const n = 400
|
||||
dur := make([]float64, n)
|
||||
first := make([]float64, n)
|
||||
var index atomic.Int64
|
||||
var failures atomic.Int64
|
||||
var wg sync.WaitGroup
|
||||
runtime.GC()
|
||||
var before, after runtime.MemStats
|
||||
runtime.ReadMemStats(&before)
|
||||
start := time.Now()
|
||||
for c := 0; c < concurrency; c++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for {
|
||||
i := int(index.Add(1)) - 1
|
||||
if i >= n {
|
||||
return
|
||||
}
|
||||
d, f, ok := request()
|
||||
dur[i] = float64(d) / 1e6
|
||||
first[i] = float64(f) / 1e6
|
||||
if !ok {
|
||||
failures.Add(1)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
elapsed := time.Since(start)
|
||||
runtime.ReadMemStats(&after)
|
||||
sort.Float64s(dur)
|
||||
sort.Float64s(first)
|
||||
var total float64
|
||||
for _, d := range dur {
|
||||
total += d
|
||||
}
|
||||
drained := elapsed
|
||||
if s != nil {
|
||||
router.Close()
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
drained = time.Since(start)
|
||||
var count int
|
||||
if err := s.db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&count); err != nil || count != n+20 {
|
||||
t.Fatalf("log count %d err %v", count, err)
|
||||
}
|
||||
}
|
||||
t.Logf("RESULT rps=%.1f drained_rps=%.1f mean_ms=%.3f p50_ms=%.3f p95_ms=%.3f p99_ms=%.3f first_p50_ms=%.3f first_p95_ms=%.3f alloc_B_req=%d errors=%d", float64(n)/elapsed.Seconds(), float64(n)/drained.Seconds(), total/n, dur[n/2], dur[n*95/100], dur[n*99/100], first[n/2], first[n*95/100], (after.TotalAlloc-before.TotalAlloc)/n, failures.Load())
|
||||
if failures.Load() != 0 {
|
||||
t.Fail()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
up.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPerformanceConversion(b *testing.B) {
|
||||
for _, size := range []int{1024, 65536} {
|
||||
for _, to := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
||||
b.Run(fmt.Sprintf("%d/%s", size, to), func(b *testing.B) {
|
||||
body := []byte(`{"model":"abc/gpt-5","max_tokens":1024,"messages":[{"role":"user","content":"` + strings.Repeat("x", size) + `"}]}`)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := protocol.ConvertRequest(protocol.Chat, to, body, "gpt-4"); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPerformanceStats(b *testing.B) {
|
||||
for _, n := range []int{1000, 10000, 100000} {
|
||||
b.Run(fmt.Sprint(n), func(b *testing.B) {
|
||||
s, err := New(filepath.Join(b.TempDir(), "stats.db"), "test-password", nil)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
_, err = s.db.Exec(`WITH RECURSIVE seq(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM seq WHERE x<?) INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) SELECT '2026-10-20T10:00:00Z','key','perf','abc','abc/gpt-5','gpt-4','openai-chat','openai-chat',200,10,0,1000,100,100,0,1,'','{}' FROM seq`, n)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
r := httptest.NewRequest("GET", "/admin/api/stats", nil)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
w := httptest.NewRecorder()
|
||||
s.stats(w, r)
|
||||
if w.Code != 200 {
|
||||
b.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func resolve(cs []channel, k apiKey, model string) (channel, string, error) {
|
||||
name, alias, ok := strings.Cut(model, "/")
|
||||
if !ok || alias == "" {
|
||||
return channel{}, "", errors.New("model must be channel/alias")
|
||||
}
|
||||
for _, c := range cs {
|
||||
if c.Name != name {
|
||||
continue
|
||||
}
|
||||
if !c.Enabled || !k.allows(c.ID) {
|
||||
return channel{}, "", errors.New("channel unavailable")
|
||||
}
|
||||
if actual, ok := c.Models[alias]; ok {
|
||||
return c, actual, nil
|
||||
}
|
||||
return channel{}, "", errors.New("model unavailable")
|
||||
}
|
||||
return channel{}, "", errors.New("channel unavailable")
|
||||
}
|
||||
func (s *Server) proxy(w http.ResponseWriter, r *http.Request) {
|
||||
kinds := map[string]protocol.Kind{"/v1/chat/completions": protocol.Chat, "/v1/responses": protocol.Responses, "/v1/messages": protocol.Anthropic}
|
||||
down, ok := kinds[r.URL.Path]
|
||||
models := r.URL.Path == "/v1/models"
|
||||
if !ok && !models {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
if (models && r.Method != "GET") || (!models && r.Method != "POST") {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
ip := s.clientIP(r)
|
||||
if !s.authAllowed(ip) {
|
||||
w.Header().Set("Retry-After", "60")
|
||||
errJSON(w, 429, "request limit exceeded")
|
||||
return
|
||||
}
|
||||
k, e := s.authenticate(r)
|
||||
if e != nil {
|
||||
s.authFail(ip)
|
||||
errJSON(w, 401, "invalid API key")
|
||||
return
|
||||
}
|
||||
release, retry := s.limiter.acquire(k.ID, time.Now())
|
||||
if release == nil {
|
||||
w.Header().Set("Retry-After", retry)
|
||||
errJSON(w, 429, "request limit exceeded")
|
||||
return
|
||||
}
|
||||
start := time.Now()
|
||||
l := requestLog{CreatedAt: start.UTC().Format(time.RFC3339Nano), KeyID: k.ID, KeyName: k.Name, Downstream: string(down), Status: 500}
|
||||
// Release the concurrency slot before the log write, which can block on a
|
||||
// full queue; rate-limited and unauthenticated requests stay unlogged.
|
||||
defer func() {
|
||||
release()
|
||||
if !models {
|
||||
l.DurationMS = time.Since(start).Milliseconds()
|
||||
s.record(l)
|
||||
}
|
||||
}()
|
||||
fail := func(status int, msg string) { l.Status = status; l.Error = msg; errJSON(w, status, msg) }
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil {
|
||||
fail(500, "database error")
|
||||
return
|
||||
}
|
||||
if models {
|
||||
ids := []string{}
|
||||
for _, c := range cs {
|
||||
if c.Enabled && k.allows(c.ID) {
|
||||
for alias := range c.Models {
|
||||
ids = append(ids, c.Name+"/"+alias)
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.Strings(ids)
|
||||
out := []any{}
|
||||
for _, id := range ids {
|
||||
out = append(out, map[string]string{"id": id, "object": "model", "owned_by": "ai-router"})
|
||||
}
|
||||
l.Status = 200
|
||||
writeJSON(w, map[string]any{"object": "list", "data": out})
|
||||
return
|
||||
}
|
||||
var b []byte
|
||||
e = readBody(w, r, maxBody, bodyReadTimeout, func() error {
|
||||
var err error
|
||||
b, err = io.ReadAll(r.Body)
|
||||
return err
|
||||
})
|
||||
if e != nil {
|
||||
var tooBig *http.MaxBytesError
|
||||
if errors.As(e, &tooBig) {
|
||||
fail(413, "request body too large")
|
||||
} else {
|
||||
fail(400, "request body unreadable")
|
||||
}
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
if json.Unmarshal(b, &req) != nil || req.Model == "" {
|
||||
fail(400, "invalid request")
|
||||
return
|
||||
}
|
||||
l.RequestModel = req.Model
|
||||
l.Stream = req.Stream
|
||||
c, actual, e := resolve(cs, k, req.Model)
|
||||
if e != nil {
|
||||
fail(403, "model or channel unavailable")
|
||||
return
|
||||
}
|
||||
l.Channel = c.Name
|
||||
l.ActualModel = actual
|
||||
l.Upstream = c.Protocol
|
||||
up := protocol.Kind(c.Protocol)
|
||||
payload, e := protocol.ConvertRequest(down, up, b, actual)
|
||||
if e != nil {
|
||||
fail(400, "unsupported request conversion")
|
||||
return
|
||||
}
|
||||
if req.Stream && up == protocol.Chat {
|
||||
var obj map[string]json.RawMessage
|
||||
if json.Unmarshal(payload, &obj) == nil {
|
||||
options := map[string]json.RawMessage{}
|
||||
if raw, ok := obj["stream_options"]; ok {
|
||||
_ = json.Unmarshal(raw, &options)
|
||||
}
|
||||
if options == nil {
|
||||
options = map[string]json.RawMessage{}
|
||||
}
|
||||
options["include_usage"] = json.RawMessage("true")
|
||||
obj["stream_options"], _ = json.Marshal(options)
|
||||
payload, _ = json.Marshal(obj)
|
||||
}
|
||||
}
|
||||
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
request, e := http.NewRequestWithContext(ctx, "POST", c.BaseURL+suffix, bytes.NewReader(payload))
|
||||
if e != nil {
|
||||
fail(502, "invalid upstream configuration")
|
||||
return
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if up == protocol.Anthropic {
|
||||
request.Header.Set("x-api-key", c.APIKey)
|
||||
request.Header.Set("anthropic-version", "2023-06-01")
|
||||
if down == up {
|
||||
if v := r.Header.Get("anthropic-version"); safeForwardedHeader(v, 32) {
|
||||
request.Header.Set("anthropic-version", v)
|
||||
}
|
||||
if v := r.Header.Get("anthropic-beta"); safeForwardedHeader(v, 256) {
|
||||
request.Header.Set("anthropic-beta", v)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
request.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
response, e := s.client.Do(request)
|
||||
if e != nil {
|
||||
fail(502, "upstream request failed")
|
||||
return
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
status := response.StatusCode
|
||||
if status < 400 {
|
||||
status = 502
|
||||
}
|
||||
fail(status, "upstream rejected request")
|
||||
return
|
||||
}
|
||||
if req.Stream {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
w.WriteHeader(response.StatusCode)
|
||||
l.Status = response.StatusCode
|
||||
streamWriter := &deadlineWriter{ResponseWriter: w}
|
||||
u, e := protocol.Stream(ctx, streamWriter, response.Body, up, down, req.Model)
|
||||
l.Usage = u
|
||||
if e != nil {
|
||||
l.Status = 502
|
||||
l.Error = "upstream stream interrupted"
|
||||
l.Usage.Known = false
|
||||
if ctx.Err() == nil {
|
||||
event := "data: {\"error\":{\"type\":\"upstream_error\",\"message\":\"upstream stream interrupted\"}}\n\n"
|
||||
if down != protocol.Chat {
|
||||
event = "event: error\n" + event
|
||||
}
|
||||
_, _ = io.WriteString(streamWriter, event)
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
b, e = io.ReadAll(io.LimitReader(response.Body, maxBody+1))
|
||||
if e != nil || len(b) > maxBody {
|
||||
fail(502, "invalid upstream response")
|
||||
return
|
||||
}
|
||||
l.Usage = protocol.ExtractUsage(up, b)
|
||||
out, e := protocol.ConvertResponse(up, down, b)
|
||||
if e != nil {
|
||||
fail(502, "unsupported response conversion")
|
||||
return
|
||||
}
|
||||
l.Status = response.StatusCode
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(response.StatusCode)
|
||||
if _, e = w.Write(out); e != nil {
|
||||
l.Status = 499
|
||||
l.Error = "client disconnected"
|
||||
}
|
||||
}
|
||||
|
||||
type deadlineWriter struct{ http.ResponseWriter }
|
||||
|
||||
// safeForwardedHeader reports whether v is a short printable token list that
|
||||
// is safe to forward upstream; anything else is dropped so clients cannot
|
||||
// smuggle arbitrary features or exotic header content past the channel
|
||||
// configuration.
|
||||
func safeForwardedHeader(v string, max int) bool {
|
||||
if len(v) == 0 || len(v) > max {
|
||||
return false
|
||||
}
|
||||
for i := 0; i < len(v); i++ {
|
||||
c := v[i]
|
||||
switch {
|
||||
case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9':
|
||||
case c == '-' || c == '_' || c == '.' || c == ',' || c == ' ':
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (w *deadlineWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
|
||||
func (w *deadlineWriter) WriteHeader(status int) {
|
||||
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||
w.ResponseWriter.WriteHeader(status)
|
||||
}
|
||||
|
||||
func (w *deadlineWriter) Write(b []byte) (int, error) {
|
||||
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||
return w.ResponseWriter.Write(b)
|
||||
}
|
||||
func (w *deadlineWriter) FlushError() error {
|
||||
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||
return http.NewResponseController(w.ResponseWriter).Flush()
|
||||
}
|
||||
func (w *deadlineWriter) Flush() { _ = w.FlushError() }
|
||||
@@ -0,0 +1,203 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const bodyReadTimeout = 30 * time.Second
|
||||
const maxLimitEntries = 10000
|
||||
|
||||
// Login rate limits: failures only, per IP and globally; successful logins
|
||||
// never consume or hit these buckets.
|
||||
const (
|
||||
loginFailuresPerMinute = 10
|
||||
loginGlobalFailuresPerMinute = 300
|
||||
)
|
||||
|
||||
// authFailuresPerMinute bounds how many failed /v1 authentications one client
|
||||
// IP may attempt per minute before further attempts are rejected before the
|
||||
// database lookup.
|
||||
const authFailuresPerMinute = 60
|
||||
|
||||
func pruneFailures(m map[string][]time.Time, now time.Time) {
|
||||
for k, ts := range m {
|
||||
if len(ts) == 0 || now.Sub(ts[len(ts)-1]) > time.Minute {
|
||||
delete(m, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// trimRecent drops entries older than one minute while reusing the backing
|
||||
// array, so a long-lived failure window cannot grow without bound.
|
||||
func trimRecent(ts []time.Time, now time.Time) []time.Time {
|
||||
i := 0
|
||||
for i < len(ts) && now.Sub(ts[i]) >= time.Minute {
|
||||
i++
|
||||
}
|
||||
if i > 0 {
|
||||
ts = append(ts[:0], ts[i:]...)
|
||||
}
|
||||
return ts
|
||||
}
|
||||
|
||||
// authAllowed reports whether the client IP may attempt authentication. Only
|
||||
// failures are counted, so legitimate traffic is unaffected.
|
||||
func (s *Server) authAllowed(ip string) bool {
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if now.Sub(s.authFailuresSeen) >= time.Minute {
|
||||
pruneFailures(s.authFailures, now)
|
||||
s.authFailuresSeen = now
|
||||
}
|
||||
recent := 0
|
||||
for _, t := range s.authFailures[ip] {
|
||||
if now.Sub(t) < time.Minute {
|
||||
recent++
|
||||
}
|
||||
}
|
||||
return recent < authFailuresPerMinute && (recent > 0 || len(s.authFailures) < maxLimitEntries)
|
||||
}
|
||||
|
||||
// authFail records one failed authentication from ip.
|
||||
func (s *Server) authFail(ip string) {
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if _, ok := s.authFailures[ip]; !ok && len(s.authFailures) >= maxLimitEntries {
|
||||
return
|
||||
}
|
||||
s.authFailures[ip] = append(s.authFailures[ip], now)
|
||||
}
|
||||
|
||||
func readBody(w http.ResponseWriter, r *http.Request, limit int64, timeout time.Duration, read func() error) error {
|
||||
controller := http.NewResponseController(w)
|
||||
_ = controller.SetReadDeadline(time.Now().Add(timeout))
|
||||
defer controller.SetReadDeadline(time.Time{})
|
||||
r.Body = http.MaxBytesReader(w, r.Body, limit)
|
||||
defer r.Body.Close()
|
||||
return read()
|
||||
}
|
||||
|
||||
// SetTrustedProxies must be called before serving requests.
|
||||
func (s *Server) SetTrustedProxies(cidrs []string) error {
|
||||
prefixes := make([]netip.Prefix, 0, len(cidrs))
|
||||
for _, cidr := range cidrs {
|
||||
p, err := netip.ParsePrefix(cidr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid trusted proxy CIDR %q: %w", cidr, err)
|
||||
}
|
||||
if p.Bits() == 0 {
|
||||
return fmt.Errorf("trusted proxy CIDR must not trust all addresses")
|
||||
}
|
||||
prefixes = append(prefixes, p)
|
||||
}
|
||||
s.trustedProxies = prefixes
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) trusted(ip netip.Addr) bool {
|
||||
for _, p := range s.trustedProxies {
|
||||
if p.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *Server) clientIP(r *http.Request) string {
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
host = r.RemoteAddr
|
||||
}
|
||||
peer, err := netip.ParseAddr(host)
|
||||
if err != nil {
|
||||
return host
|
||||
}
|
||||
peer = peer.Unmap()
|
||||
if !s.trusted(peer) {
|
||||
return peer.String()
|
||||
}
|
||||
parts := strings.Split(strings.Join(r.Header.Values("X-Forwarded-For"), ","), ",")
|
||||
ips := make([]netip.Addr, len(parts))
|
||||
for i, part := range parts {
|
||||
ip, err := netip.ParseAddr(strings.TrimSpace(part))
|
||||
if err != nil || ip.Zone() != "" {
|
||||
return peer.String()
|
||||
}
|
||||
ips[i] = ip.Unmap()
|
||||
}
|
||||
for i := len(ips) - 1; i >= 0; i-- {
|
||||
if !s.trusted(ips[i]) {
|
||||
return ips[i].String()
|
||||
}
|
||||
}
|
||||
return peer.String()
|
||||
}
|
||||
|
||||
type limitEntry struct {
|
||||
window time.Time
|
||||
requests, active int
|
||||
}
|
||||
type requestLimiter struct {
|
||||
mu sync.Mutex
|
||||
rpm, perKey, global, active int
|
||||
entries map[string]*limitEntry
|
||||
nextCleanup time.Time
|
||||
}
|
||||
|
||||
func newRequestLimiter(rpm, perKey, global int) *requestLimiter {
|
||||
return &requestLimiter{rpm: rpm, perKey: perKey, global: global, entries: make(map[string]*limitEntry)}
|
||||
}
|
||||
|
||||
// SetLimits must be called before serving requests.
|
||||
func (s *Server) SetLimits(rpm, perKey, global int) error {
|
||||
if rpm <= 0 || perKey <= 0 || global <= 0 {
|
||||
return fmt.Errorf("limits must be positive integers")
|
||||
}
|
||||
s.limiter = newRequestLimiter(rpm, perKey, global)
|
||||
return nil
|
||||
}
|
||||
func (l *requestLimiter) acquire(key string, now time.Time) (func(), string) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
if !now.Before(l.nextCleanup) {
|
||||
for k, e := range l.entries {
|
||||
if e.active == 0 && now.Sub(e.window) >= time.Minute {
|
||||
delete(l.entries, k)
|
||||
}
|
||||
}
|
||||
l.nextCleanup = now.Add(time.Minute)
|
||||
}
|
||||
e := l.entries[key]
|
||||
if e == nil {
|
||||
if len(l.entries) >= maxLimitEntries {
|
||||
return nil, "60"
|
||||
}
|
||||
e = &limitEntry{window: now}
|
||||
l.entries[key] = e
|
||||
}
|
||||
if now.Sub(e.window) >= time.Minute {
|
||||
e.window = now
|
||||
e.requests = 0
|
||||
}
|
||||
if e.requests >= l.rpm {
|
||||
seconds := int((time.Minute - now.Sub(e.window) + time.Second - 1) / time.Second)
|
||||
return nil, strconv.Itoa(seconds)
|
||||
}
|
||||
if e.active >= l.perKey || l.active >= l.global {
|
||||
return nil, "1"
|
||||
}
|
||||
e.requests++
|
||||
e.active++
|
||||
l.active++
|
||||
var once sync.Once
|
||||
return func() { once.Do(func() { l.mu.Lock(); defer l.mu.Unlock(); e.active--; l.active-- }) }, ""
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func TestSecurityStatsSQLSemantics(t *testing.T) {
|
||||
s := testServer(t)
|
||||
now := time.Now().UTC().Format(time.RFC3339Nano)
|
||||
for _, l := range []requestLog{
|
||||
{CreatedAt: now, KeyID: "a", KeyName: "old", Status: 200, Usage: protocol.Usage{Known: true, Input: 8}},
|
||||
{CreatedAt: now, KeyID: "a", KeyName: "new", Status: 200, Error: "failed", Usage: protocol.Usage{Raw: []byte("{}"), Input: 100}},
|
||||
{CreatedAt: now, KeyID: "b", Status: 400},
|
||||
{CreatedAt: now, KeyID: "c", Status: 204, Usage: protocol.Usage{Known: true}},
|
||||
} {
|
||||
s.record(l)
|
||||
}
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
query string
|
||||
requests, success, unknown int64
|
||||
input *int64
|
||||
}{
|
||||
{"", 4, 2, 2, new(int64(8))},
|
||||
{"?key_id=b", 1, 0, 1, nil},
|
||||
{"?key_id=c", 1, 1, 0, new(int64(0))},
|
||||
{"?key_id=missing", 0, 0, 0, nil},
|
||||
} {
|
||||
w := httptest.NewRecorder()
|
||||
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats"+tc.query, nil))
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var got struct {
|
||||
Summary summary
|
||||
ByKey []struct{ Name string } `json:"by_key"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Summary.Requests != tc.requests || got.Summary.Success != tc.success || got.Summary.Unknown != tc.unknown {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
if (got.Summary.Input == nil) != (tc.input == nil) || tc.input != nil && *got.Summary.Input != *tc.input {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
if tc.query == "" && got.ByKey[0].Name != "new" {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLogFieldsBounded(t *testing.T) {
|
||||
s := testServer(t)
|
||||
huge := strings.Repeat("界", 10000)
|
||||
s.record(requestLog{RequestModel: huge, ActualModel: huge, Channel: huge, KeyID: huge, KeyName: huge, Upstream: huge, Downstream: huge, Error: huge, Usage: protocol.Usage{Known: true, Input: 12, Raw: []byte(huge)}})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var model, actual, channel, key, name, up, down, message, raw string
|
||||
var input int64
|
||||
if err := s.db.QueryRow("SELECT request_model,actual_model,channel,key_id,key_name,upstream,downstream,error,raw_usage,input_tokens FROM logs").Scan(&model, &actual, &channel, &key, &name, &up, &down, &message, &raw, &input); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, field := range []string{model, actual, channel, key, name, up, down} {
|
||||
if len(field) > maxLogFieldBytes || !utf8.ValidString(field) {
|
||||
t.Fatal("unbounded or invalid field")
|
||||
}
|
||||
}
|
||||
if len(message) > maxLogErrorBytes || !utf8.ValidString(message) || len(raw) > maxLogUsageBytes || !json.Valid([]byte(raw)) || input != 12 {
|
||||
t.Fatal("invalid bounded log")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLogRetentionOptInAndBatch(t *testing.T) {
|
||||
s := testServer(t)
|
||||
old := time.Now().Add(-91 * 24 * time.Hour).Format(time.RFC3339Nano)
|
||||
for i := 0; i < logCleanupBatch+3; i++ {
|
||||
s.record(requestLog{CreatedAt: old})
|
||||
}
|
||||
s.record(requestLog{CreatedAt: time.Now().Format(time.RFC3339Nano)})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s.logRetentionDays.Load() != 0 {
|
||||
t.Fatal("retention enabled by default")
|
||||
}
|
||||
if err := s.purgeExpiredLogs(context.Background(), 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != logCleanupBatch+4 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
s.SetLogRetention(90)
|
||||
if s.logRetentionDays.Load() != 90 {
|
||||
t.Fatal("retention not set")
|
||||
}
|
||||
if err := s.purgeExpiredLogs(context.Background(), 90); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != 4 {
|
||||
t.Fatal("unbounded purge", n)
|
||||
}
|
||||
if err := s.purgeExpiredLogs(context.Background(), 90); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
s.SetLogRetention(0)
|
||||
if s.logRetentionDays.Load() != 0 {
|
||||
t.Fatal("retention not disabled")
|
||||
}
|
||||
s.SetLogRetention(90)
|
||||
if err := s.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsGroupLimitAndCancellation(t *testing.T) {
|
||||
s := testServer(t)
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 10001; i++ {
|
||||
if err := insertLog(context.Background(), tx, requestLog{CreatedAt: "2026-01-01T00:00:00Z", KeyID: fmt.Sprint(i)}); err != nil {
|
||||
tx.Rollback()
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil))
|
||||
if w.Code != 422 {
|
||||
t.Fatal("unbounded group result", w.Code)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
w = httptest.NewRecorder()
|
||||
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil).WithContext(ctx))
|
||||
if w.Code == 200 {
|
||||
t.Fatal("cancelled query succeeded")
|
||||
}
|
||||
var day string
|
||||
if err := s.db.QueryRow("SELECT router_log_day(NULL, 'UTC')").Scan(&day); err == nil {
|
||||
t.Fatal("invalid timestamp accepted")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestClientIPTrust(t *testing.T) {
|
||||
s := &Server{}
|
||||
if err := s.SetTrustedProxies([]string{"10.0.0.0/8", "::1/128"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetTrustedProxies([]string{"bad"}); err == nil {
|
||||
t.Fatal("invalid CIDR accepted")
|
||||
}
|
||||
for _, tt := range []struct{ peer, xff, want string }{
|
||||
{"192.0.2.1:42", "198.51.100.1", "192.0.2.1"},
|
||||
{"10.0.0.1:42", "198.51.100.1, 10.0.0.2", "198.51.100.1"},
|
||||
{"10.0.0.1:42", "198.51.100.1, 192.0.2.2", "192.0.2.2"},
|
||||
{"10.0.0.1:42", "bad, 192.0.2.2", "10.0.0.1"},
|
||||
{"10.0.0.1:42", "", "10.0.0.1"},
|
||||
{"10.0.0.1:42", "10.0.0.2", "10.0.0.1"},
|
||||
{"[::1]:42", "2001:db8::1", "2001:db8::1"},
|
||||
} {
|
||||
r := httptest.NewRequest("POST", "/", nil)
|
||||
r.RemoteAddr = tt.peer
|
||||
r.Header.Set("X-Forwarded-For", tt.xff)
|
||||
r.Header.Set("X-Real-IP", "203.0.113.1")
|
||||
if got := s.clientIP(r); got != tt.want {
|
||||
t.Errorf("%+v: got %s", tt, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginProxyIsolation(t *testing.T) {
|
||||
s, err := New(":memory:", "password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
login := func(peer, forwarded string) int {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"wrong"}`))
|
||||
r.RemoteAddr = peer
|
||||
r.Header.Set("X-Forwarded-For", forwarded)
|
||||
w := httptest.NewRecorder()
|
||||
s.login(w, r)
|
||||
return w.Code
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.1"); got != 401 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.1"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.2"); got != 401 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
login("192.0.2.3:1", fmt.Sprintf("198.51.100.%d", i))
|
||||
}
|
||||
if got := login("192.0.2.3:1", "203.0.113.1"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBodyDeadline(t *testing.T) {
|
||||
result := make(chan error, 1)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
|
||||
result <- err
|
||||
w.WriteHeader(400)
|
||||
}))
|
||||
defer srv.Close()
|
||||
conn, err := net.Dial("tcp", srv.Listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 20\r\n\r\nx")
|
||||
select {
|
||||
case err := <-result:
|
||||
if err == nil {
|
||||
t.Fatal("slow body accepted")
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("body deadline did not fire")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBodyDeadlineCleared(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
|
||||
if err != nil {
|
||||
http.Error(w, "read", 400)
|
||||
return
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
http.Error(w, "canceled", 500)
|
||||
return
|
||||
default:
|
||||
}
|
||||
fmt.Fprint(w, "data: ok\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
conn, err := net.Dial("tcp", srv.Listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
conn.SetDeadline(time.Now().Add(2 * time.Second))
|
||||
fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 1\r\n\r\nx")
|
||||
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatal(resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginBodyLimit(t *testing.T) {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+strings.Repeat("x", 16<<10)+`"}`))
|
||||
var v any
|
||||
if decode(httptest.NewRecorder(), r, &v) == nil {
|
||||
t.Fatal("oversize login accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterReleaseAndCleanup(t *testing.T) {
|
||||
now := time.Now()
|
||||
l := newRequestLimiter(2, 1, 2)
|
||||
a, _ := l.acquire("a", now)
|
||||
if a == nil {
|
||||
t.Fatal("denied")
|
||||
}
|
||||
if release, _ := l.acquire("a", now); release != nil {
|
||||
t.Fatal("per-key concurrency bypass")
|
||||
}
|
||||
b, _ := l.acquire("b", now)
|
||||
if b == nil {
|
||||
t.Fatal("denied")
|
||||
}
|
||||
if release, _ := l.acquire("c", now); release != nil {
|
||||
t.Fatal("global concurrency bypass")
|
||||
}
|
||||
a()
|
||||
a()
|
||||
b()
|
||||
a, _ = l.acquire("a", now)
|
||||
if a == nil {
|
||||
t.Fatal("not released")
|
||||
}
|
||||
a()
|
||||
if release, retry := l.acquire("a", now); release != nil || retry != "60" {
|
||||
t.Fatal("rpm bypass")
|
||||
}
|
||||
c, _ := l.acquire("c", now.Add(time.Minute))
|
||||
if c == nil {
|
||||
t.Fatal("window not reset")
|
||||
}
|
||||
c()
|
||||
if len(l.entries) != 1 {
|
||||
t.Fatalf("stale entries: %d", len(l.entries))
|
||||
}
|
||||
for i := 0; i < maxLimitEntries; i++ {
|
||||
l.entries[fmt.Sprint(i)] = &limitEntry{window: now.Add(time.Minute)}
|
||||
}
|
||||
if release, _ := l.acquire("overflow", now.Add(time.Minute)); release != nil {
|
||||
t.Fatal("map cap bypass")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterConcurrent(t *testing.T) {
|
||||
l := newRequestLimiter(10000, 8, 64)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if release, _ := l.acquire("key", time.Now()); release != nil {
|
||||
release()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if l.active != 0 || l.entries["key"].active != 0 {
|
||||
t.Fatal("permit leak")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyLimitRejectsBeforeBodyAndLogging(t *testing.T) {
|
||||
s, err := New(":memory:", "test-password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetLimits(1, 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||
key := integrationCreateKey(t, s, cookie, "limited", ch)
|
||||
request := func(body io.Reader) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", body)
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
if w := request(strings.NewReader("{}")); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if s.limiter.active != 0 {
|
||||
t.Fatal("permit leaked after invalid request")
|
||||
}
|
||||
body := &observedReader{}
|
||||
w := request(body)
|
||||
if w.Code != 429 || w.Header().Get("Retry-After") == "" || body.read {
|
||||
t.Fatalf("status=%d read=%t", w.Code, body.read)
|
||||
}
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var count int
|
||||
if err := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&count); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Fatalf("rejected request logged: %d logs", count)
|
||||
}
|
||||
}
|
||||
|
||||
type observedReader struct{ read bool }
|
||||
|
||||
func (r *observedReader) Read(p []byte) (int, error) { r.read = true; return 0, io.EOF }
|
||||
|
||||
func TestLoginFailuresOnlyAndGlobalCap(t *testing.T) {
|
||||
s, err := New(":memory:", "password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
login := func(peer, forwarded, password string) int {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+password+`"}`))
|
||||
r.RemoteAddr = peer
|
||||
if forwarded != "" {
|
||||
r.Header.Set("X-Forwarded-For", forwarded)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.login(w, r)
|
||||
return w.Code
|
||||
}
|
||||
for i := 0; i < 15; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.1", "password"); got != 200 {
|
||||
t.Fatalf("correct login %d rejected: %d", i, got)
|
||||
}
|
||||
}
|
||||
for i := 0; i < loginFailuresPerMinute; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 401 {
|
||||
t.Fatalf("failure %d: %d", i, got)
|
||||
}
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
// The per-IP lockout is uniform: it must not confirm candidate passwords.
|
||||
if got := login("10.0.0.1:1", "192.0.2.2", "password"); got != 429 {
|
||||
t.Fatal("lockout leaks password validity:", got)
|
||||
}
|
||||
// A successful login clears its IP bucket, so fat-fingers do not accumulate.
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "password"); got != 200 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for i := 0; i < loginFailuresPerMinute; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
|
||||
t.Fatalf("post-clear failure %d: %d", i, got)
|
||||
}
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for i := 0; i < loginGlobalFailuresPerMinute; i++ {
|
||||
login("10.0.0.1:1", fmt.Sprintf("198.18.%d.%d", i/256, i%256), "wrong")
|
||||
}
|
||||
if got := login("10.0.0.1:1", "203.0.113.1", "wrong"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
// The global cap fails closed too: aggregate brute force locks every
|
||||
// login out without revealing candidate-password validity.
|
||||
if got := login("10.0.0.1:1", "203.0.113.1", "password"); got != 429 {
|
||||
t.Fatal("global lockout leaks password validity:", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthFailureGate(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||
key := integrationCreateKey(t, s, cookie, "gated", ch)
|
||||
request := func(auth string) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(integrationRequest("openai-chat", "abc/gpt-5", false)))
|
||||
r.Header.Set("Authorization", "Bearer "+auth)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
if w := request(key.raw); w.Code != 502 {
|
||||
t.Fatal("valid key setup:", w.Code, w.Body.String())
|
||||
}
|
||||
for i := 0; i < authFailuresPerMinute; i++ {
|
||||
if w := request("wrong"); w.Code != 401 {
|
||||
t.Fatalf("%d: %d", i, w.Code)
|
||||
}
|
||||
}
|
||||
if w := request("wrong"); w.Code != 429 || w.Header().Get("Retry-After") == "" {
|
||||
t.Fatal(w.Code, w.Header())
|
||||
}
|
||||
// The gate applies before the key lookup: a valid key from the same
|
||||
// (attacker-shared) IP is throttled the same way.
|
||||
if w := request(key.raw); w.Code != 429 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyBodyTooLarge(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||
key := integrationCreateKey(t, s, cookie, "limit", ch)
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"abc/gpt-5","messages":[],"pad":"`+strings.Repeat("x", maxBody)+`"}`))
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 413 || !strings.Contains(w.Body.String(), "too large") {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"io/fs"
|
||||
_ "modernc.org/sqlite"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Server struct {
|
||||
logRetentionDays atomic.Int64
|
||||
trustedProxies []netip.Prefix
|
||||
limiter *requestLimiter
|
||||
db *sql.DB
|
||||
password [32]byte
|
||||
assets fs.FS
|
||||
mu sync.Mutex
|
||||
sessions map[string]time.Time
|
||||
attempts map[string][]time.Time
|
||||
authFailures map[string][]time.Time
|
||||
authFailuresSeen time.Time
|
||||
loginFailures []time.Time
|
||||
client *http.Client
|
||||
publicOrigin string
|
||||
lifecycleMu sync.Mutex
|
||||
closing bool
|
||||
active sync.WaitGroup
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
logMu sync.RWMutex
|
||||
logsClosed bool
|
||||
logQueue chan logEntry
|
||||
logDone chan struct{}
|
||||
logErr error
|
||||
}
|
||||
|
||||
const maxBody = 16 << 20
|
||||
|
||||
func New(p, password string, a fs.FS) (*Server, error) {
|
||||
if password == "" {
|
||||
return nil, errors.New("empty admin password")
|
||||
}
|
||||
if p != ":memory:" {
|
||||
if e := os.MkdirAll(filepath.Dir(p), 0700); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
}
|
||||
d, e := sql.Open("sqlite", p)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
d.SetMaxOpenConns(1)
|
||||
schema := `PRAGMA busy_timeout=5000; CREATE TABLE IF NOT EXISTS channels(id TEXT PRIMARY KEY,name TEXT NOT NULL UNIQUE,base_url TEXT NOT NULL,protocol TEXT NOT NULL,enabled INTEGER NOT NULL,passthrough INTEGER NOT NULL,models TEXT NOT NULL,api_key TEXT NOT NULL); CREATE TABLE IF NOT EXISTS keys(id TEXT PRIMARY KEY,name TEXT NOT NULL,prefix TEXT NOT NULL,hash TEXT NOT NULL UNIQUE,enabled INTEGER NOT NULL,all_channels INTEGER NOT NULL,channels TEXT NOT NULL,expires_at TEXT NOT NULL,created_at TEXT NOT NULL,revoked INTEGER NOT NULL DEFAULT 0); CREATE TABLE IF NOT EXISTS logs(id INTEGER PRIMARY KEY AUTOINCREMENT,created_at TEXT NOT NULL,key_id TEXT NOT NULL,key_name TEXT NOT NULL,channel TEXT NOT NULL,request_model TEXT NOT NULL,actual_model TEXT NOT NULL,downstream TEXT NOT NULL,upstream TEXT NOT NULL,status INTEGER NOT NULL,duration_ms INTEGER NOT NULL,stream INTEGER NOT NULL,input_tokens INTEGER,output_tokens INTEGER,cache_read_tokens INTEGER,cache_write_tokens INTEGER,usage_known INTEGER NOT NULL,error TEXT NOT NULL,raw_usage TEXT NOT NULL); CREATE INDEX IF NOT EXISTS logs_time ON logs(created_at); CREATE INDEX IF NOT EXISTS logs_key ON logs(key_id);`
|
||||
if _, e = d.Exec(schema); e != nil {
|
||||
d.Close()
|
||||
return nil, e
|
||||
}
|
||||
if p != ":memory:" {
|
||||
if e = os.Chmod(p, 0600); e != nil {
|
||||
d.Close()
|
||||
return nil, e
|
||||
}
|
||||
}
|
||||
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||
// Upstreams must not be able to turn compressed streams into proxy-side
|
||||
// decompression work; the router asks for identity encoding.
|
||||
tr.DisableCompression = true
|
||||
tr.ResponseHeaderTimeout = 60 * time.Second
|
||||
tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
c, e := (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, addr)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
return &idleConn{Conn: c}, nil
|
||||
}
|
||||
s := &Server{db: d, password: sha256.Sum256([]byte(password)), assets: a, sessions: map[string]time.Time{}, attempts: map[string][]time.Time{}, authFailures: map[string][]time.Time{}, client: &http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, logQueue: make(chan logEntry, 1024), logDone: make(chan struct{})}
|
||||
go s.writeLogs()
|
||||
s.limiter = newRequestLimiter(120, 8, 64)
|
||||
return s, nil
|
||||
}
|
||||
func (s *Server) Close() error {
|
||||
s.closeOnce.Do(func() {
|
||||
s.lifecycleMu.Lock()
|
||||
s.closing = true
|
||||
s.lifecycleMu.Unlock()
|
||||
s.active.Wait()
|
||||
s.logMu.Lock()
|
||||
s.logsClosed = true
|
||||
close(s.logQueue)
|
||||
s.logMu.Unlock()
|
||||
<-s.logDone
|
||||
s.client.CloseIdleConnections()
|
||||
s.closeErr = errors.Join(s.logErr, s.db.Close())
|
||||
})
|
||||
return s.closeErr
|
||||
}
|
||||
func writeJSON(w http.ResponseWriter, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
func errJSON(w http.ResponseWriter, n int, m string) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(n)
|
||||
json.NewEncoder(w).Encode(map[string]string{"error": m})
|
||||
}
|
||||
func decode(w http.ResponseWriter, r *http.Request, v any) error {
|
||||
limit := int64(maxBody)
|
||||
if r.URL.Path == "/admin/api/login" {
|
||||
limit = 16 << 10
|
||||
}
|
||||
return readBody(w, r, limit, bodyReadTimeout, func() error {
|
||||
d := json.NewDecoder(r.Body)
|
||||
if e := d.Decode(v); e != nil {
|
||||
return e
|
||||
}
|
||||
var extra any
|
||||
if d.Decode(&extra) != io.EOF {
|
||||
return errors.New("trailing JSON")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
func randomID() string {
|
||||
b := make([]byte, 32)
|
||||
if _, e := rand.Read(b); e != nil {
|
||||
panic(e)
|
||||
}
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
func hash(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
|
||||
func (s *Server) session(r *http.Request) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
c, e := r.Cookie("ai_router_session")
|
||||
if e != nil {
|
||||
return false
|
||||
}
|
||||
t, ok := s.sessions[c.Value]
|
||||
return ok && time.Now().Before(t)
|
||||
}
|
||||
func (s *Server) guard(r *http.Request) bool {
|
||||
if r.Header.Get("X-Requested-With") != "ai-router" {
|
||||
return false
|
||||
}
|
||||
o, e := url.Parse(r.Header.Get("Origin"))
|
||||
if s.publicOrigin != "" {
|
||||
// Origins are scheme/host case-insensitive; a mixed-case deployment
|
||||
// value must not lock the admin out.
|
||||
return e == nil && strings.EqualFold(r.Header.Get("Origin"), s.publicOrigin)
|
||||
}
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
return e == nil && o.Scheme == scheme && strings.EqualFold(o.Host, r.Host) && o.User == nil && o.RawQuery == "" && o.Fragment == "" && o.Path == ""
|
||||
}
|
||||
func (s *Server) secureCookie() bool {
|
||||
return len(s.publicOrigin) >= 8 && strings.EqualFold(s.publicOrigin[:8], "https://")
|
||||
}
|
||||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
w = &deadlineWriter{ResponseWriter: w}
|
||||
_ = http.NewResponseController(w).SetReadDeadline(time.Now().Add(bodyReadTimeout))
|
||||
s.lifecycleMu.Lock()
|
||||
if s.closing {
|
||||
s.lifecycleMu.Unlock()
|
||||
errJSON(w, http.StatusServiceUnavailable, "server shutting down")
|
||||
return
|
||||
}
|
||||
s.active.Add(1)
|
||||
s.lifecycleMu.Unlock()
|
||||
defer s.active.Done()
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Referrer-Policy", "no-referrer")
|
||||
w.Header().Set("X-Frame-Options", "DENY")
|
||||
w.Header().Set("Content-Security-Policy", "default-src 'none'; base-uri 'none'; frame-ancestors 'none'; form-action 'none'; connect-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:")
|
||||
if strings.HasPrefix(r.URL.Path, "/admin/api/") || r.URL.Path == "/admin/api" {
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
s.admin(w, r)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/v1/") {
|
||||
s.proxy(w, r)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/api/") {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if s.assets == nil || (r.Method != "GET" && r.Method != "HEAD") {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
p := strings.TrimPrefix(r.URL.Path, "/")
|
||||
if p == "" {
|
||||
p = "index.html"
|
||||
}
|
||||
if st, e := fs.Stat(s.assets, p); e != nil || st.IsDir() {
|
||||
p = "index.html"
|
||||
}
|
||||
http.ServeFileFS(w, r, s.assets, p)
|
||||
}
|
||||
func (s *Server) admin(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != "GET" && !s.guard(r) {
|
||||
errJSON(w, 403, "forbidden")
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/admin/api/login" && r.Method == "POST" {
|
||||
s.login(w, r)
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/admin/api/session" && r.Method == "GET" {
|
||||
writeJSON(w, map[string]bool{"authenticated": s.session(r)})
|
||||
return
|
||||
}
|
||||
if !s.session(r) {
|
||||
errJSON(w, 401, "unauthorized")
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/admin/api/logout" && r.Method == "POST" {
|
||||
c, _ := r.Cookie("ai_router_session")
|
||||
s.mu.Lock()
|
||||
delete(s.sessions, c.Value)
|
||||
s.mu.Unlock()
|
||||
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", MaxAge: -1, Path: "/", HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
||||
writeJSON(w, map[string]bool{"authenticated": false})
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case r.URL.Path == "/admin/api/models/test":
|
||||
s.testModel(w, r)
|
||||
case r.URL.Path == "/admin/api/channels":
|
||||
s.channels(w, r)
|
||||
case r.URL.Path == "/admin/api/channels/discover-models":
|
||||
s.discoverModels(w, r)
|
||||
case strings.HasPrefix(r.URL.Path, "/admin/api/channels/"):
|
||||
s.channelID(w, r)
|
||||
case r.URL.Path == "/admin/api/keys":
|
||||
s.keys(w, r)
|
||||
case strings.HasPrefix(r.URL.Path, "/admin/api/keys/"):
|
||||
s.keyID(w, r)
|
||||
case r.URL.Path == "/admin/api/logs" && r.Method == "GET":
|
||||
s.logs(w, r)
|
||||
case r.URL.Path == "/admin/api/stats" && r.Method == "GET":
|
||||
s.stats(w, r)
|
||||
default:
|
||||
errJSON(w, 404, "not found")
|
||||
}
|
||||
}
|
||||
// login gates on failure counts BEFORE verifying the password. While an IP is
|
||||
// locked — or the global failure budget is spent — every attempt, including a
|
||||
// correct password, gets the same 429 without a comparison, so responses never
|
||||
// reveal whether a candidate password is right and brute force stays capped
|
||||
// at loginFailuresPerMinute per IP and loginGlobalFailuresPerMinute overall.
|
||||
// Successful logins are not failures: their gate slot is released by clearing
|
||||
// the IP bucket, and only wrong passwords feed the global window.
|
||||
func (s *Server) login(w http.ResponseWriter, r *http.Request) {
|
||||
ip := s.clientIP(r)
|
||||
now := time.Now()
|
||||
var x struct {
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if decode(w, r, &x) != nil {
|
||||
errJSON(w, 400, "invalid request")
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
pruneFailures(s.attempts, now)
|
||||
s.loginFailures = trimRecent(s.loginFailures, now)
|
||||
recent := []time.Time{}
|
||||
for _, t := range s.attempts[ip] {
|
||||
if now.Sub(t) < time.Minute {
|
||||
recent = append(recent, t)
|
||||
}
|
||||
}
|
||||
blocked := len(recent) >= loginFailuresPerMinute || len(s.loginFailures) >= loginGlobalFailuresPerMinute || (len(recent) == 0 && len(s.attempts) >= 10000)
|
||||
if !blocked {
|
||||
// Reserve the gate slot atomically with the check so a concurrent
|
||||
// burst cannot slip past the limit; a success erases the bucket below.
|
||||
s.attempts[ip] = append(recent, now)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if blocked {
|
||||
w.Header().Set("Retry-After", "60")
|
||||
errJSON(w, 429, "too many attempts")
|
||||
return
|
||||
}
|
||||
h := sha256.Sum256([]byte(x.Password))
|
||||
if subtle.ConstantTimeCompare(h[:], s.password[:]) != 1 {
|
||||
s.mu.Lock()
|
||||
s.loginFailures = append(s.loginFailures, now)
|
||||
s.mu.Unlock()
|
||||
errJSON(w, 401, "invalid credentials")
|
||||
return
|
||||
}
|
||||
v := randomID()
|
||||
s.mu.Lock()
|
||||
delete(s.attempts, ip)
|
||||
for k, t := range s.sessions {
|
||||
if now.After(t) {
|
||||
delete(s.sessions, k)
|
||||
}
|
||||
}
|
||||
if c, e := r.Cookie("ai_router_session"); e == nil {
|
||||
delete(s.sessions, c.Value)
|
||||
}
|
||||
s.sessions[v] = now.Add(24 * time.Hour)
|
||||
s.mu.Unlock()
|
||||
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", Value: v, Path: "/", MaxAge: 86400, HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
||||
writeJSON(w, map[string]bool{"authenticated": true})
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"testing/fstest"
|
||||
"time"
|
||||
)
|
||||
|
||||
func testServer(t *testing.T) *Server {
|
||||
t.Helper()
|
||||
s, e := New(":memory:", "test-password", fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("SPA")}})
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
t.Cleanup(func() { s.Close() })
|
||||
return s
|
||||
}
|
||||
func call(s *Server, m, p, b string, c *http.Cookie) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest(m, p, strings.NewReader(b))
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
if c != nil {
|
||||
r.AddCookie(c)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
func loginCookie(t *testing.T, s *Server) *http.Cookie {
|
||||
t.Helper()
|
||||
w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
return w.Result().Cookies()[0]
|
||||
}
|
||||
func TestAuthentication(t *testing.T) {
|
||||
s := testServer(t)
|
||||
for _, origin := range []string{"", "http://evil.test", "https://example.com"} {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||
r.Header.Set("Origin", origin)
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(origin, w.Code)
|
||||
}
|
||||
}
|
||||
c := loginCookie(t, s)
|
||||
if !c.HttpOnly || c.SameSite != http.SameSiteStrictMode {
|
||||
t.Fatal(c)
|
||||
}
|
||||
if w := call(s, "GET", "/admin/api/channels", "", nil); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := call(s, "POST", "/admin/api/logout", "", c); w.Code != 200 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := call(s, "GET", "/admin/api/channels", "", c); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := call(s, "GET", "/admin/api/session", "", nil); w.Code != 200 || w.Header().Get("Content-Security-Policy") == "" || w.Header().Get("X-Frame-Options") != "DENY" || w.Header().Get("Referrer-Policy") != "no-referrer" {
|
||||
t.Fatal("security headers missing", w.Code, w.Header())
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil)
|
||||
}
|
||||
if w := call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil); w.Code != 429 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
// While locked, a correct password gets the same 429: the lockout must
|
||||
// not become an oracle that confirms candidate passwords.
|
||||
if w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil); w.Code != 429 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
// The lockout is per IP; a fresh address still logs in.
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
r.RemoteAddr = "192.0.2.9:1"
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 200 {
|
||||
t.Fatal("fresh IP rejected:", w.Code)
|
||||
}
|
||||
}
|
||||
func TestMappingPermissions(t *testing.T) {
|
||||
cs := []channel{{ID: "stable", Name: "renamed", Enabled: true, Models: map[string]string{"alias": "vendor/real", "vendor/real": "wrong"}}}
|
||||
k := apiKey{Channels: []string{"stable"}}
|
||||
_, actual, e := resolve(cs, k, "renamed/alias")
|
||||
if e != nil || actual != "vendor/real" {
|
||||
t.Fatal(actual, e)
|
||||
}
|
||||
if _, _, e = resolve(cs, apiKey{Channels: []string{"renamed"}}, "renamed/alias"); e == nil {
|
||||
t.Fatal("name granted access")
|
||||
}
|
||||
if _, _, e = resolve(cs, k, "renamed/vendor/unknown"); e == nil {
|
||||
t.Fatal("unknown allowed")
|
||||
}
|
||||
cs[0].Enabled = false
|
||||
if _, _, e = resolve(cs, k, "renamed/alias"); e == nil {
|
||||
t.Fatal("disabled allowed")
|
||||
}
|
||||
}
|
||||
func TestLegacyPassthroughIgnored(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Error("unmapped model reached upstream")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
}))
|
||||
defer up.Close()
|
||||
body := map[string]any{
|
||||
"name": "legacy", "base_url": up.URL, "protocol": "openai-chat",
|
||||
"enabled": true, "passthrough": true, "models": map[string]string{"alias": "vendor/real"},
|
||||
}
|
||||
save := func() channel {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
var c channel
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &c); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return c
|
||||
}
|
||||
c := save()
|
||||
var stored bool
|
||||
if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
|
||||
t.Fatal("new channel enabled passthrough", stored, err)
|
||||
}
|
||||
body["id"] = c.ID
|
||||
save()
|
||||
if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
|
||||
t.Fatal("update enabled passthrough", stored, err)
|
||||
}
|
||||
if _, err := s.db.Exec("UPDATE channels SET passthrough=1 WHERE id=?", c.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cs, err := s.loadChannels()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, actual, err := resolve(cs, apiKey{AllChannels: true}, "legacy/alias"); err != nil || actual != "vendor/real" {
|
||||
t.Fatal(actual, err)
|
||||
}
|
||||
w := call(s, "GET", "/admin/api/channels", "", cookie)
|
||||
if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
w = call(s, "POST", "/admin/api/keys", `{"name":"client","all_channels":true}`, cookie)
|
||||
var key map[string]string
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &key); err != nil || w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String(), err)
|
||||
}
|
||||
for _, path := range []string{"/v1/chat/completions", "/v1/responses", "/v1/messages"} {
|
||||
for _, model := range []string{"legacy/unknown", "legacy/vendor/real"} {
|
||||
t.Run(path+"/"+model, func(t *testing.T) {
|
||||
r := httptest.NewRequest("POST", path, strings.NewReader(`{"model":"`+model+`","messages":[]}`))
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
save()
|
||||
cs, err = s.loadChannels()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := resolve(cs, apiKey{AllChannels: true}, "legacy/unknown"); err == nil {
|
||||
t.Fatal("updated legacy channel allowed unknown model")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyRevocationAndProxy(t *testing.T) {
|
||||
s := testServer(t)
|
||||
c := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/chat/completions" || r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
||||
t.Error(r.URL.Path, r.Header)
|
||||
}
|
||||
var b map[string]any
|
||||
json.NewDecoder(r.Body).Decode(&b)
|
||||
if b["model"] != "vendor/real" || b["custom"] != "keep" {
|
||||
t.Error(b)
|
||||
}
|
||||
writeJSON(w, map[string]any{"choices": []any{}, "usage": map[string]int{"prompt_tokens": 7, "completion_tokens": 3}})
|
||||
}))
|
||||
defer up.Close()
|
||||
b, _ := json.Marshal(channel{Name: "demo", BaseURL: up.URL + "/v1", Protocol: "openai-chat", Enabled: true, Models: map[string]string{"alias": "vendor/real"}, APIKey: "upstream-secret"})
|
||||
w := call(s, "POST", "/admin/api/channels", string(b), c)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var ch channel
|
||||
json.Unmarshal(w.Body.Bytes(), &ch)
|
||||
if strings.Contains(w.Body.String(), "upstream-secret") {
|
||||
t.Fatal("secret leak")
|
||||
}
|
||||
w = call(s, "POST", "/admin/api/keys", `{"name":"client","channels":["`+ch.ID+`"]}`, c)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var key map[string]string
|
||||
json.Unmarshal(w.Body.Bytes(), &key)
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"demo/alias","messages":[],"custom":"keep"}`))
|
||||
r.Header.Set("x-api-key", key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
var id, stored string
|
||||
if e := s.db.QueryRow("SELECT id,hash FROM keys").Scan(&id, &stored); e != nil || stored == key["key"] {
|
||||
t.Fatal(e, stored)
|
||||
}
|
||||
w = call(s, "DELETE", "/admin/api/keys/"+id, "", c)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
w = call(s, "PATCH", "/admin/api/keys/"+id, `{"enabled":true}`, c)
|
||||
if w.Code != 404 {
|
||||
t.Fatal("revived revoked key", w.Code)
|
||||
}
|
||||
r = httptest.NewRequest("GET", "/v1/models", nil)
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
var n int
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.db.QueryRow("SELECT count(*) FROM logs").Scan(&n)
|
||||
if n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
func TestStatsUnknownAndFilters(t *testing.T) {
|
||||
s := testServer(t)
|
||||
c := loginCookie(t, s)
|
||||
now := time.Now().UTC().Format(time.RFC3339Nano)
|
||||
s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 200, Usage: protocol.Usage{Input: 8, Output: 2, Known: true}})
|
||||
s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 400})
|
||||
s.record(requestLog{CreatedAt: now, KeyID: "b", Status: 200})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := call(s, "GET", "/admin/api/stats?key_id=a", "", c)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var out struct {
|
||||
Summary summary `json:"summary"`
|
||||
}
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &out); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if out.Summary.Requests != 2 || out.Summary.Success != 1 || out.Summary.Unknown != 1 || out.Summary.Input == nil || *out.Summary.Input != 8 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
w = call(s, "GET", "/admin/api/logs?key_id=b", "", c)
|
||||
if !strings.Contains(w.Body.String(), `"input_tokens":null`) {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
w = call(s, "GET", "/admin/api/stats?from=invalid", "", c)
|
||||
if w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
w = call(s, "GET", "/admin/api/unknown", "", c)
|
||||
if w.Code != 404 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
w = call(s, "GET", "/some/spa/path", "", nil)
|
||||
if w.Code != 200 || w.Body.String() != "SPA" {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublicOriginCaseInsensitive(t *testing.T) {
|
||||
s := testServer(t)
|
||||
if e := s.SetPublicOrigin("https://Router.Example.com"); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||
r.Header.Set("Origin", "https://router.example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 200 || !w.Result().Cookies()[0].Secure {
|
||||
t.Fatal(w.Code, w.Header())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublicOriginAndRedirect(t *testing.T) {
|
||||
s := testServer(t)
|
||||
if e := s.SetPublicOrigin("https://router.example"); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||
r.Header.Set("Origin", "https://router.example")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 200 || !w.Result().Cookies()[0].Secure {
|
||||
t.Fatal(w.Code, w.Header())
|
||||
}
|
||||
reached := false
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached = true }))
|
||||
defer target.Close()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
||||
defer upstream.Close()
|
||||
req, _ := http.NewRequest("POST", upstream.URL, strings.NewReader("body"))
|
||||
req.Header.Set("Authorization", "Bearer private")
|
||||
res, e := s.client.Do(req)
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
res.Body.Close()
|
||||
if reached || res.StatusCode != 307 {
|
||||
t.Fatal("redirect followed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"context"
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
_ "time/tzdata"
|
||||
|
||||
"modernc.org/sqlite"
|
||||
)
|
||||
|
||||
func init() {
|
||||
var locationMu sync.Mutex
|
||||
var cachedZone string
|
||||
var cachedLocation *time.Location
|
||||
if err := sqlite.RegisterDeterministicScalarFunction("router_log_day", 2, func(_ *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
|
||||
stampText, ok := args[0].(string)
|
||||
zone, zoneOK := args[1].(string)
|
||||
if !ok || !zoneOK {
|
||||
return nil, fmt.Errorf("invalid log timestamp or timezone")
|
||||
}
|
||||
stamp, err := time.Parse(time.RFC3339Nano, stampText)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
locationMu.Lock()
|
||||
defer locationMu.Unlock()
|
||||
if cachedLocation == nil || cachedZone != zone {
|
||||
loc, err := time.LoadLocation(zone)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cachedZone, cachedLocation = zone, loc
|
||||
}
|
||||
return stamp.In(cachedLocation).Format("2006-01-02"), nil
|
||||
}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
type requestLog struct {
|
||||
ID int64 `json:"id"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
KeyID string `json:"key_id"`
|
||||
KeyName string `json:"key_name"`
|
||||
Channel string `json:"channel"`
|
||||
RequestModel string `json:"request_model"`
|
||||
ActualModel string `json:"actual_model"`
|
||||
Downstream string `json:"downstream"`
|
||||
Upstream string `json:"upstream"`
|
||||
Status int `json:"status"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
Stream bool `json:"stream"`
|
||||
Input *int64 `json:"input_tokens"`
|
||||
Output *int64 `json:"output_tokens"`
|
||||
CacheRead *int64 `json:"cache_read_tokens"`
|
||||
CacheWrite *int64 `json:"cache_write_tokens"`
|
||||
UsageKnown bool `json:"usage_known"`
|
||||
Error string `json:"error"`
|
||||
RawUsage json.RawMessage `json:"raw_usage"`
|
||||
Usage protocol.Usage `json:"-"`
|
||||
}
|
||||
|
||||
func insertLog(ctx context.Context, tx *sql.Tx, l requestLog) error {
|
||||
l = boundedLog(l)
|
||||
var in, out, cr, cw any
|
||||
if l.Usage.Known || len(l.Usage.Raw) > 0 {
|
||||
in = l.Usage.Input
|
||||
out = l.Usage.Output
|
||||
cr = l.Usage.CacheRead
|
||||
cw = l.Usage.CacheWrite
|
||||
}
|
||||
_, err := tx.ExecContext(ctx, "INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", l.CreatedAt, l.KeyID, l.KeyName, l.Channel, l.RequestModel, l.ActualModel, l.Downstream, l.Upstream, l.Status, l.DurationMS, l.Stream, in, out, cr, cw, l.Usage.Known, l.Error, string(l.Usage.Raw))
|
||||
return err
|
||||
}
|
||||
func filter(r *http.Request) (string, []any, error) {
|
||||
parts := []string{"1=1"}
|
||||
args := []any{}
|
||||
for _, name := range []string{"channel", "model"} {
|
||||
if v := r.URL.Query().Get(name); v != "" {
|
||||
col := name
|
||||
if name == "model" {
|
||||
col = "actual_model"
|
||||
}
|
||||
parts = append(parts, col+"=?")
|
||||
args = append(args, v)
|
||||
}
|
||||
}
|
||||
if k := r.URL.Query().Get("key_id"); k != "" {
|
||||
parts = append(parts, "key_id=?")
|
||||
args = append(args, k)
|
||||
}
|
||||
var from, to time.Time
|
||||
for _, n := range []string{"from", "to"} {
|
||||
if v := r.URL.Query().Get(n); v != "" {
|
||||
t, e := time.Parse(time.RFC3339Nano, v)
|
||||
if e != nil {
|
||||
return "", nil, e
|
||||
}
|
||||
if n == "from" {
|
||||
from = t
|
||||
parts = append(parts, "julianday(created_at)>=julianday(?)")
|
||||
} else {
|
||||
to = t
|
||||
parts = append(parts, "julianday(created_at)<=julianday(?)")
|
||||
}
|
||||
args = append(args, t.UTC().Format(time.RFC3339Nano))
|
||||
}
|
||||
}
|
||||
if !from.IsZero() && !to.IsZero() && from.After(to) {
|
||||
return "", nil, sql.ErrNoRows
|
||||
}
|
||||
return strings.Join(parts, " AND "), args, nil
|
||||
}
|
||||
func (s *Server) readLogs(r *http.Request, limited bool) ([]requestLog, error) {
|
||||
where, args, e := filter(r)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
q := "SELECT id,created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage FROM logs WHERE " + where + " ORDER BY id DESC"
|
||||
if limited {
|
||||
q += " LIMIT 200"
|
||||
}
|
||||
rows, e := s.db.Query(q, args...)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []requestLog{}
|
||||
for rows.Next() {
|
||||
var l requestLog
|
||||
var raw string
|
||||
if e = rows.Scan(&l.ID, &l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream, &l.Status, &l.DurationMS, &l.Stream, &l.Input, &l.Output, &l.CacheRead, &l.CacheWrite, &l.UsageKnown, &l.Error, &raw); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if json.Valid([]byte(raw)) {
|
||||
l.RawUsage = json.RawMessage(raw)
|
||||
}
|
||||
out = append(out, l)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
func (s *Server) logs(w http.ResponseWriter, r *http.Request) {
|
||||
if _, _, e := filter(r); e != nil {
|
||||
errJSON(w, 400, "invalid time filter")
|
||||
return
|
||||
}
|
||||
ls, e := s.readLogs(r, true)
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, ls)
|
||||
}
|
||||
|
||||
type summary struct {
|
||||
Requests int64 `json:"requests"`
|
||||
Success int64 `json:"success"`
|
||||
Failed int64 `json:"failed"`
|
||||
Input *int64 `json:"input_tokens"`
|
||||
Output *int64 `json:"output_tokens"`
|
||||
CacheRead *int64 `json:"cache_read_tokens"`
|
||||
CacheWrite *int64 `json:"cache_write_tokens"`
|
||||
Unknown int64 `json:"unknown_usage"`
|
||||
}
|
||||
|
||||
func (a *summary) add(l requestLog) {
|
||||
a.Requests++
|
||||
if l.Status >= 200 && l.Status < 300 && l.Error == "" {
|
||||
a.Success++
|
||||
} else {
|
||||
a.Failed++
|
||||
}
|
||||
if !l.UsageKnown {
|
||||
a.Unknown++
|
||||
return
|
||||
}
|
||||
if a.Input == nil {
|
||||
a.Input = new(int64)
|
||||
a.Output = new(int64)
|
||||
a.CacheRead = new(int64)
|
||||
a.CacheWrite = new(int64)
|
||||
}
|
||||
if l.Input != nil {
|
||||
*a.Input += *l.Input
|
||||
}
|
||||
if l.Output != nil {
|
||||
*a.Output += *l.Output
|
||||
}
|
||||
if l.CacheRead != nil {
|
||||
*a.CacheRead += *l.CacheRead
|
||||
}
|
||||
if l.CacheWrite != nil {
|
||||
*a.CacheWrite += *l.CacheWrite
|
||||
}
|
||||
}
|
||||
func (s *Server) stats(w http.ResponseWriter, r *http.Request) {
|
||||
zone := r.URL.Query().Get("timezone")
|
||||
if zone == "" {
|
||||
zone = "UTC"
|
||||
}
|
||||
_, err := time.LoadLocation(zone)
|
||||
if err != nil || zone == "Local" {
|
||||
errJSON(w, 400, "invalid timezone")
|
||||
return
|
||||
}
|
||||
where, args, e := filter(r)
|
||||
if e != nil {
|
||||
errJSON(w, 400, "invalid time filter")
|
||||
return
|
||||
}
|
||||
const aggregates = `COUNT(*), COALESCE(SUM(status>=200 AND status<300 AND error=''),0),
|
||||
COALESCE(SUM(NOT(status>=200 AND status<300 AND error='')),0),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(input_tokens,0) END),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(output_tokens,0) END),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(cache_read_tokens,0) END),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(cache_write_tokens,0) END),
|
||||
COALESCE(SUM(NOT usage_known),0), MAX(id)`
|
||||
// NOT MATERIALIZED avoids copying the full log payload into a temporary table.
|
||||
query := `WITH filtered AS NOT MATERIALIZED (SELECT * FROM logs WHERE ` + where + `)
|
||||
SELECT 'total','','',` + aggregates + ` FROM filtered
|
||||
UNION ALL SELECT 'key',key_id,key_name,` + aggregates + ` FROM filtered GROUP BY key_id
|
||||
UNION ALL SELECT 'day',router_log_day(created_at,?), '',` + aggregates + ` FROM filtered GROUP BY router_log_day(created_at,?)
|
||||
UNION ALL SELECT 'model',channel,actual_model,` + aggregates + ` FROM filtered GROUP BY channel,actual_model`
|
||||
args = append(args, zone, zone)
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
|
||||
defer cancel()
|
||||
rows, e := s.db.QueryContext(ctx, query, args...)
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
type byKey struct {
|
||||
KeyID string `json:"key_id"`
|
||||
Name string `json:"name"`
|
||||
summary
|
||||
}
|
||||
type daily struct {
|
||||
Day string `json:"day"`
|
||||
summary
|
||||
}
|
||||
type byModel struct {
|
||||
Channel string `json:"channel"`
|
||||
Model string `json:"model"`
|
||||
summary
|
||||
}
|
||||
total := summary{}
|
||||
ks := map[string]*byKey{}
|
||||
ds := map[string]*daily{}
|
||||
ms := map[[2]string]*byModel{}
|
||||
groups := 0
|
||||
for rows.Next() {
|
||||
groups++
|
||||
if groups > 10000 {
|
||||
errJSON(w, 422, "too many statistics groups; narrow the time range")
|
||||
return
|
||||
}
|
||||
var kind, a, b string
|
||||
var v summary
|
||||
var latest sql.NullInt64
|
||||
if err := rows.Scan(&kind, &a, &b, &v.Requests, &v.Success, &v.Failed, &v.Input, &v.Output, &v.CacheRead, &v.CacheWrite, &v.Unknown, &latest); err != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
switch kind {
|
||||
case "total":
|
||||
total = v
|
||||
case "key":
|
||||
// SQLite's sole MAX selects key_name from the latest ID, not timestamp.
|
||||
ks[a] = &byKey{KeyID: a, Name: b, summary: v}
|
||||
case "day":
|
||||
ds[a] = &daily{Day: a, summary: v}
|
||||
case "model":
|
||||
ms[[2]string{a, b}] = &byModel{Channel: a, Model: b, summary: v}
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
keys := []*byKey{}
|
||||
days := []*daily{}
|
||||
mods := []*byModel{}
|
||||
for _, v := range ks {
|
||||
keys = append(keys, v)
|
||||
}
|
||||
for _, v := range ds {
|
||||
days = append(days, v)
|
||||
}
|
||||
for _, v := range ms {
|
||||
mods = append(mods, v)
|
||||
}
|
||||
sort.Slice(keys, func(i, j int) bool { return keys[i].KeyID < keys[j].KeyID })
|
||||
sort.Slice(days, func(i, j int) bool { return days[i].Day < days[j].Day })
|
||||
sort.Slice(mods, func(i, j int) bool { return mods[i].Channel+mods[i].Model < mods[j].Channel+mods[j].Model })
|
||||
writeJSON(w, map[string]any{"summary": total, "by_key": keys, "daily": days, "by_model": mods, "timezone": zone})
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestStatsBrowserTimezone(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, zone string
|
||||
stamps []string
|
||||
want map[string]int
|
||||
}{
|
||||
{"Shanghai midnight", "Asia/Shanghai", []string{"2026-10-20T15:59:59Z", "2026-10-20T16:00:00Z"}, map[string]int{"2026-10-20": 1, "2026-10-21": 1}},
|
||||
{"UTC default", "", []string{"2026-10-20T15:59:59Z", "2026-10-20T16:00:00Z"}, map[string]int{"2026-10-20": 2}},
|
||||
{"New York DST spring", "America/New_York", []string{"2025-03-09T04:59:59Z", "2025-03-09T05:00:00Z", "2025-03-10T03:59:59Z", "2025-03-10T04:00:00Z"}, map[string]int{"2025-03-08": 1, "2025-03-09": 2, "2025-03-10": 1}},
|
||||
{"New York DST fall", "America/New_York", []string{"2025-11-02T03:59:59Z", "2025-11-02T04:00:00Z", "2025-11-03T04:59:59Z", "2025-11-03T05:00:00Z"}, map[string]int{"2025-11-01": 1, "2025-11-02": 2, "2025-11-03": 1}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
for _, stamp := range tc.stamps {
|
||||
s.record(requestLog{CreatedAt: stamp, KeyID: "one", Status: 200, Usage: protocol.Usage{Known: true, Input: 10}})
|
||||
}
|
||||
s.record(requestLog{CreatedAt: tc.stamps[0], KeyID: "other", Status: 200})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := call(s, "GET", "/admin/api/stats?key_id=one&timezone="+url.QueryEscape(tc.zone), "", cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
var result struct {
|
||||
Timezone string `json:"timezone"`
|
||||
Summary summary `json:"summary"`
|
||||
Daily []struct {
|
||||
Day string `json:"day"`
|
||||
Requests int `json:"requests"`
|
||||
} `json:"daily"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := map[string]int{}
|
||||
for _, d := range result.Daily {
|
||||
got[d.Day] = d.Requests
|
||||
}
|
||||
if !reflect.DeepEqual(got, tc.want) {
|
||||
t.Fatalf("got %v want %v", got, tc.want)
|
||||
}
|
||||
if result.Summary.Requests != int64(len(tc.stamps)) || result.Summary.Input == nil || *result.Summary.Input != int64(10*len(tc.stamps)) {
|
||||
t.Fatal(result.Summary)
|
||||
}
|
||||
var stored string
|
||||
if err := s.db.QueryRow("SELECT created_at FROM logs ORDER BY id LIMIT 1").Scan(&stored); err != nil || stored != tc.stamps[0] {
|
||||
t.Fatal("UTC storage changed", stored, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsRejectsInvalidTimezone(t *testing.T) {
|
||||
s := testServer(t)
|
||||
c := loginCookie(t, s)
|
||||
for _, zone := range []string{"Local", "not/a-zone", "../../etc/passwd"} {
|
||||
w := call(s, "GET", "/admin/api/stats?timezone="+url.QueryEscape(zone), "", c)
|
||||
if w.Code != 400 {
|
||||
t.Fatal(zone, w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net"
|
||||
"net/url"
|
||||
"time"
|
||||
)
|
||||
|
||||
type idleConn struct{ net.Conn }
|
||||
|
||||
func (c *idleConn) Read(b []byte) (int, error) {
|
||||
if e := c.SetReadDeadline(time.Now().Add(2 * time.Minute)); e != nil {
|
||||
return 0, e
|
||||
}
|
||||
return c.Conn.Read(b)
|
||||
}
|
||||
func (c *idleConn) Write(b []byte) (int, error) {
|
||||
if e := c.SetWriteDeadline(time.Now().Add(30 * time.Second)); e != nil {
|
||||
return 0, e
|
||||
}
|
||||
return c.Conn.Write(b)
|
||||
}
|
||||
|
||||
// SetPublicOrigin must be called before serving requests; forwarded headers are not trusted.
|
||||
func (s *Server) SetPublicOrigin(origin string) error {
|
||||
u, e := url.Parse(origin)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" {
|
||||
return errors.New("public origin must be an http(s) origin")
|
||||
}
|
||||
s.publicOrigin = origin
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type observedDeadlineWriter struct {
|
||||
*httptest.ResponseRecorder
|
||||
writes, reads int
|
||||
deadline time.Time
|
||||
}
|
||||
|
||||
func (w *observedDeadlineWriter) SetWriteDeadline(d time.Time) error {
|
||||
w.writes++
|
||||
w.deadline = d
|
||||
return nil
|
||||
}
|
||||
func (w *observedDeadlineWriter) SetReadDeadline(d time.Time) error { w.reads++; return nil }
|
||||
|
||||
func TestResponseWriteDeadlines(t *testing.T) {
|
||||
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
||||
w := &deadlineWriter{ResponseWriter: raw}
|
||||
w.WriteHeader(200)
|
||||
if raw.writes != 1 || !raw.deadline.After(time.Now()) {
|
||||
t.Fatal("headers lack deadline")
|
||||
}
|
||||
if _, err := w.Write([]byte("body")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := w.FlushError(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if raw.writes != 3 {
|
||||
t.Fatal("write/flush must refresh deadline", raw.writes)
|
||||
}
|
||||
if err := http.NewResponseController(w).SetReadDeadline(time.Time{}); err != nil || raw.reads != 1 {
|
||||
t.Fatal("wrapper must preserve read deadline control", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNonStreamResponseHasWriteDeadline(t *testing.T) {
|
||||
s := testServer(t)
|
||||
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
||||
s.ServeHTTP(raw, httptest.NewRequest("GET", "/admin/api/session", nil))
|
||||
if raw.Code != 200 || raw.writes == 0 || !raw.deadline.After(time.Now()) {
|
||||
t.Fatal("nonstream response lacks write deadline")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user