From 45b9d778bb6093a7ea2f0e86bfc7ab99887ce82c Mon Sep 17 00:00:00 2001 From: Vasil Date: Wed, 28 Oct 2026 01:19:01 +0100 Subject: [PATCH 1/2] feat(http): Anthropic Messages API at /v1/messages for Claude Code and the Anthropic SDKs --- CHANGELOG.md | 9 + README.md | 63 +++ ROADMAP.md | 11 - internal/grpcapi/handler.go | 8 +- internal/grpcapi/handler_test.go | 2 +- internal/server/anthropic.go | 660 ++++++++++++++++++++++++++++++ internal/server/anthropic_test.go | 285 +++++++++++++ internal/server/openai.go | 124 +++--- plugin/skills/integrate/SKILL.md | 6 +- 9 files changed, 1096 insertions(+), 72 deletions(-) create mode 100644 internal/server/anthropic.go create mode 100644 internal/server/anthropic_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 023795d..80ad6cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ All notable changes to KTMS are listed here. The format follows ## [Unreleased] +### Added +- Anthropic Messages API (experimental): `POST /v1/messages`, so Claude Code + (`ANTHROPIC_BASE_URL`) and the Anthropic SDKs run through Keystone with a KTMS key in + `x-api-key` or `Authorization`, `model: "auto"` or a registry model id. It shares routing, + fallback, limits, guardrails, cache, `X-KTMS-*` headers and usage events with chat completions, + answers in the Messages format whichever vendor's model serves it, streams the Messages SSE + events, and returns errors in the Messages shape. `POST /v1/messages/count_tokens` returns an + estimate. + ## [1.4.0] - 2026-10-27 ### Added diff --git a/README.md b/README.md index cfc1b46..81526c3 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ vendor is down, enforces rate and spend limits, and records usage. - [Command-line tool](#command-line-tool) - [Calling Keystone](#calling-keystone) - [OpenAI-compatible HTTP API](#openai-compatible-http-api) +- [Anthropic Messages API](#anthropic-messages-api) - [MCP gateway](#mcp-gateway) - [Concepts](#concepts) - [Error contract](#error-contract) @@ -1399,6 +1400,66 @@ Errors use OpenAI's shape, `{"error": {"message", "type", "param", "code"}}`, an OpenAI SDKs retry 429 and 5xx automatically; for `throttled:*` and `vendor_error` you may want to disable that (`max_retries=0`) and handle them yourself. +## Anthropic Messages API + +`POST /v1/messages` serves the [Anthropic Messages API](https://docs.anthropic.com/en/api/messages) +on `HTTP_ADDR`, so Claude Code and the Anthropic SDKs run through Keystone. It is the same pipeline +as `/v1/chat/completions`: the same keys (scope `generate`), rules, fallback, limits, guardrails, +cache, `X-KTMS-*` headers and usage events. The model that answers may be any vendor's, and the +answer is still a Messages response. The Messages API is [experimental](#experimental) in 1.x. + +Claude Code, with a KTMS key: + +```bash +export ANTHROPIC_BASE_URL=http://localhost:8080 +export ANTHROPIC_AUTH_TOKEN=ktms_ # or ANTHROPIC_API_KEY +export ANTHROPIC_MODEL=auto # or a registry model id +export ANTHROPIC_DEFAULT_HAIKU_MODEL=auto:low # Claude Code's background calls +claude +``` + +Claude Code warns once that it does not recognize the model `auto`; that is expected. It asks for +up to 32,000-64,000 output tokens, which Keystone passes on as `max_tokens`; if a model your rules +pick caps output lower, set `CLAUDE_CODE_MAX_OUTPUT_TOKENS` to fit it. + +The Anthropic SDKs need only the base URL and the key: + +```python +import anthropic + +client = anthropic.Anthropic(base_url="http://localhost:8080", api_key="ktms_") +reply = client.messages.create(model="auto", max_tokens=1024, + messages=[{"role": "user", "content": "Say hi in three words."}]) +print(reply.content[0].text, reply.model) # `model` is the one that answered +``` + +| Request field | Support | +| --- | --- | +| `model` | As on chat completions: `auto`, `auto:`, `auto:`, or a registry model id. | +| `system` | A string or `text` blocks, joined into the system instruction. | +| `messages` | `user` and `assistant` turns with `text`, `image` and `document` blocks (`base64` or `https` `url` sources; `text` and `content` documents become text), `tool_use` and `tool_result` (text and images). `thinking`, `redacted_thinking` and server-tool blocks in assistant turns are dropped: another model cannot continue them. | +| `tools`, `tool_choice` | Client tools with `input_schema` (the same JSON Schema subset as `response_format`); `tool_choice` `auto`, `any`, `tool` and `none`. Server tools such as `web_search` are rejected with `400`. | +| `max_tokens` | Required, the output ceiling. | +| `stream` | SSE events `message_start`, `ping`, `content_block_start`, `content_block_delta` (`text_delta`, `input_json_delta`), `content_block_stop`, `message_delta` and `message_stop`. `message_start` carries an input token estimate; `message_delta` carries the vendor's counts. | +| `metadata.user_id` | `attributes["user_id"]`, as `user` on chat completions; other `metadata` keys are attributes, and `action_id` is the usage action label. | +| `enable_cache`, `placement` | Extensions, as on chat completions; so are the `X-KTMS-Cache` and `X-KTMS-Placement` request headers (set them in Claude Code with `ANTHROPIC_CUSTOM_HEADERS`). | +| `temperature`, `top_p`, `top_k`, `stop_sequences`, `thinking`, `service_tier` | Accepted and ignored. | + +The key goes in `x-api-key` or `Authorization: Bearer`. The response's `stop_reason` is `end_turn`, +`tool_use`, `max_tokens` or `refusal`, and `usage` splits `cache_read_input_tokens` and +`cache_creation_input_tokens` out of `input_tokens`. Tool call ids are the vendor's, or `toolu_…` +when the vendor gives none. + +`POST /v1/messages/count_tokens` answers with an estimate, not a tokenizer's count: about four +characters per token, tool definitions included, with the `X-KTMS-Token-Count: estimate` header. + +Errors use the Messages shape, `{"type": "error", "error": {"type", "message"}}`, with the HTTP +statuses of [the table above](#errors): `invalid_request_error` (400), `authentication_error` +(401), `permission_error` (403), `not_found_error` (404), `rate_limit_error` (429, with +`Retry-After`), `overloaded_error` (503), `timeout_error` (504) and `api_error` for the rest. The +machine-readable reason, such as `throttled:rpm`, is the `message`. An error after the stream has +started arrives as an `error` event. + ## MCP gateway Keystone is also a gateway for MCP servers. Claude Code, Cursor and other MCP clients reach the @@ -1601,6 +1662,8 @@ These work and are supported, but may change in a minor release. The change is t - **`JudgeStream` and `EvaluateRollouts`**, and `POST /v1/decisions/rollouts` over HTTP. - **The MCP gateway:** `/mcp`, the `MCPServerSpec` RPCs, `allowed_mcp`, the `mcp_servers` and `mcp_credentials` file sections, the `ktms_mcp_*` metrics and `ToolCall` in usage events. +- **The Anthropic Messages API:** `POST /v1/messages`, `POST /v1/messages/count_tokens` and the + `X-KTMS-Token-Count` header. ### Compatibility with ktms-admin and light-tokenmeter diff --git a/ROADMAP.md b/ROADMAP.md index 5791485..a580ea2 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -16,19 +16,8 @@ Every 1.x release is additive: new RPCs, fields, endpoints and settings. Anythin ## Now -### Agent tools: MCP gateway - -- A registry of MCP servers (Streamable HTTP and SSE), proxied through Keystone at `/mcp`. -- One endpoint that aggregates several servers, with tools namespaced per server. -- Access to servers and individual tools per client key and team, like `allowed_models` today. -- Upstream credentials stored encrypted and injected by the gateway, never handed to clients. -- Tool calls in light-tokenmeter and Prometheus, and limits per tool. -- Managed the same way as models: declarative YAML, Admin API, CLI and the console. - ### Coding agents and API -- Anthropic Messages API at `/v1/messages`, so Claude Code runs through KTMS with - `ANTHROPIC_BASE_URL` and `model: "auto"`. - OpenAI Responses API at `/v1/responses` for Codex CLI and the Agents SDK. - `temperature`, `top_p`, `stop` and `seed` passed through to the chosen model. diff --git a/internal/grpcapi/handler.go b/internal/grpcapi/handler.go index c32f72d..4e135d0 100644 --- a/internal/grpcapi/handler.go +++ b/internal/grpcapi/handler.go @@ -154,7 +154,7 @@ func (s *ProxyServer) GenerateText(ctx context.Context, req *pb.GenerateTextRequ if plan.shadowEval { shadowEvaluator = s.makeShadowEvaluator(id, chat, req.GetAttributes()) } - estimate := estimateTokens(req) + estimate := EstimateTokens(req) settleKey, err := s.admitKey(ctx, id, estimate) if err != nil { return nil, err @@ -297,7 +297,7 @@ func (s *ProxyServer) GenerateTextStream(req *pb.GenerateTextStreamRequest, stre if plan.shadowEval { shadowEvaluator = s.makeShadowEvaluator(id, chat, inner.GetAttributes()) } - estimate := estimateTokens(inner) + estimate := EstimateTokens(inner) settleKey, err := s.admitKey(stream.Context(), id, estimate) if err != nil { return err @@ -662,9 +662,9 @@ func buildChat(req *pb.GenerateTextRequest, id *apikeys.Identity, ruleName, effo return chat } -// estimateTokens is the budget reservation guess: total runes / 4 plus heuristic +// EstimateTokens is the budget reservation guess: total runes / 4 plus heuristic // allowances for media parts. Settled against vendor-reported actuals after the call. -func estimateTokens(req *pb.GenerateTextRequest) int64 { +func EstimateTokens(req *pb.GenerateTextRequest) int64 { runes := utf8.RuneCountInString(req.GetBaseSystemInstruction()) + utf8.RuneCountInString(req.GetRequestInstruction()) var extraTokens int64 diff --git a/internal/grpcapi/handler_test.go b/internal/grpcapi/handler_test.go index 36339eb..20671a2 100644 --- a/internal/grpcapi/handler_test.go +++ b/internal/grpcapi/handler_test.go @@ -548,7 +548,7 @@ func TestEstimateTokens(t *testing.T) { BaseSystemInstruction: "abcd", // 4 runes Messages: []*pb.ChatMessage{{Text: "abcdefgh"}}, // 8 runes } - if got := estimateTokens(req); got != 4 { // 12/4 + 1 + if got := EstimateTokens(req); got != 4 { // 12/4 + 1 t.Fatalf("estimate = %d, want 4", got) } } diff --git a/internal/server/anthropic.go b/internal/server/anthropic.go new file mode 100644 index 0000000..437ec8a --- /dev/null +++ b/internal/server/anthropic.go @@ -0,0 +1,660 @@ +package server + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "strings" + "unicode/utf8" + + "github.com/rs/zerolog/log" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + pb "github.com/ai-process/ktms/gen/ktms/v1" + "github.com/ai-process/ktms/internal/grpcapi" + "github.com/ai-process/ktms/internal/safefetch" +) + +// The Anthropic Messages API runs on the same GenerateText pipeline as chat +// completions: routing, fallback, limits, guardrails, cache and usage are +// shared, and whatever model answers is rendered in the Messages format. + +type messagesRequest struct { + Model string `json:"model"` + System json.RawMessage `json:"system"` + Messages []anthropicMessage `json:"messages"` + MaxTokens int32 `json:"max_tokens"` + Stream bool `json:"stream"` + Tools []anthropicTool `json:"tools"` + ToolChoice *anthropicToolChoice `json:"tool_choice"` + Metadata map[string]string `json:"metadata"` + // Proxy extensions, as on chat completions. + EnableCache *bool `json:"enable_cache,omitempty"` + Placement string `json:"placement,omitempty"` +} + +type anthropicMessage struct { + Role string `json:"role"` + Content json.RawMessage `json:"content"` +} + +type anthropicBlock struct { + Type string `json:"type"` + Text string `json:"text"` + Source *anthropicSource `json:"source"` + ID string `json:"id"` + Name string `json:"name"` + Input json.RawMessage `json:"input"` + ToolUseID string `json:"tool_use_id"` + Content json.RawMessage `json:"content"` + IsError bool `json:"is_error"` + Title string `json:"title"` +} + +type anthropicSource struct { + Type string `json:"type"` + MediaType string `json:"media_type"` + Data string `json:"data"` + URL string `json:"url"` + Content json.RawMessage `json:"content"` +} + +type anthropicTool struct { + Type string `json:"type"` + Name string `json:"name"` + Description string `json:"description"` + InputSchema json.RawMessage `json:"input_schema"` +} + +type anthropicToolChoice struct { + Type string `json:"type"` + Name string `json:"name"` +} + +// anthropicBlocks reads content that is either a string or a list of blocks. +func anthropicBlocks(raw json.RawMessage, what string) ([]anthropicBlock, error) { + if len(raw) == 0 || string(raw) == "null" { + return nil, nil + } + var s string + if json.Unmarshal(raw, &s) == nil { + return []anthropicBlock{{Type: "text", Text: s}}, nil + } + var blocks []anthropicBlock + if err := json.Unmarshal(raw, &blocks); err != nil { + return nil, invalid("%s must be a string or a list of content blocks", what) + } + return blocks, nil +} + +// skippedBlocks are assistant blocks another model can neither read nor +// continue: thinking is signed by the model that wrote it, and server tools +// are not forwarded. +var skippedBlocks = map[string]bool{ + "thinking": true, "redacted_thinking": true, "server_tool_use": true, "web_search_tool_result": true, +} + +func (req *messagesRequest) toProto(ctx context.Context, headerPlacement string) (*pb.GenerateTextRequest, error) { + if len(req.Messages) == 0 { + return nil, invalid("messages: at least one message is required") + } + placement := req.Placement + if headerPlacement != "" { + placement = headerPlacement + } + dec, err := route(req.Model, "", placement) + if err != nil { + return nil, err + } + out := &pb.GenerateTextRequest{ + Effort: dec.Effort, + ModelOverride: dec.ModelOverride, + MaxOutputTokens: req.MaxTokens, + Placement: dec.Placement, + AutoRoute: dec.IsAuto, + EffortSpecified: dec.EffortSpecified, + } + out.Attributes, out.ActionId = attributes(req.Metadata, "") + + system, err := anthropicBlocks(req.System, "system") + if err != nil { + return nil, err + } + var texts []string + for _, b := range system { + if b.Type != "text" { + return nil, invalid("system: block type %q is not supported, only text", b.Type) + } + texts = append(texts, b.Text) + } + out.BaseSystemInstruction = strings.Join(texts, "\n\n") + + for _, t := range req.Tools { + if t.Type != "" && t.Type != "custom" { + return nil, invalid("tools: server tool %q is not supported; KTMS forwards client tools only", t.Type) + } + if t.Name == "" { + return nil, invalid("tools: name is required") + } + tool := &pb.Tool{Type: "function", Function: &pb.FunctionDefinition{Name: t.Name, Description: t.Description}} + if len(t.InputSchema) > 0 && string(t.InputSchema) != "null" { + if tool.Function.Parameters, err = responseSchemaFromJSON(t.InputSchema); err != nil { + return nil, invalid("tools: %s: invalid input_schema: %v", t.Name, err) + } + } + out.Tools = append(out.Tools, tool) + } + if tc := req.ToolChoice; tc != nil { + switch tc.Type { + case "auto": + out.ToolChoice = &pb.ToolChoice{Mode: pb.ToolChoice_MODE_AUTO} + case "any": + out.ToolChoice = &pb.ToolChoice{Mode: pb.ToolChoice_MODE_REQUIRED} + case "none": + out.ToolChoice = &pb.ToolChoice{Mode: pb.ToolChoice_MODE_NONE} + case "tool": + if tc.Name == "" { + return nil, invalid("tool_choice: name is required for type tool") + } + out.ToolChoice = &pb.ToolChoice{Mode: pb.ToolChoice_MODE_SPECIFIC, SpecificFunctionName: tc.Name} + default: + return nil, invalid("tool_choice: type %q is not supported", tc.Type) + } + } + + for i, m := range req.Messages { + blocks, err := anthropicBlocks(m.Content, fmt.Sprintf("messages[%d].content", i)) + if err != nil { + return nil, err + } + var msgs []*pb.ChatMessage + switch m.Role { + case "user": + msgs, err = userMessages(ctx, blocks) + case "assistant": + msgs, err = assistantMessage(blocks) + default: + err = invalid("role %q is not supported, use user or assistant", m.Role) + } + if err != nil { + return nil, prefixInvalid(fmt.Sprintf("messages[%d]", i), err) + } + out.Messages = append(out.Messages, msgs...) + } + return out, nil +} + +func prefixInvalid(prefix string, err error) error { + if st, ok := status.FromError(err); ok && st.Code() == codes.InvalidArgument { + return invalid("%s: %s", prefix, st.Message()) + } + return err +} + +// userMessages turns a user turn into its tool results, each a tool message, +// then one user message with the rest. Images in tool results join that user +// message: tool messages carry text only. +func userMessages(ctx context.Context, blocks []anthropicBlock) ([]*pb.ChatMessage, error) { + var out []*pb.ChatMessage + var texts []string + var parts []*pb.ContentPart + addText := func(s string) { + texts = append(texts, s) + parts = append(parts, &pb.ContentPart{Part: &pb.ContentPart_Text{Text: &pb.TextPart{Text: s}}}) + } + for _, b := range blocks { + switch b.Type { + case "text": + addText(b.Text) + case "image": + p, err := imagePart(ctx, b.Source) + if err != nil { + return nil, err + } + parts = append(parts, p) + case "document": + p, text, err := documentPart(ctx, b) + if err != nil { + return nil, err + } + if p != nil { + parts = append(parts, p) + } else { + addText(text) + } + case "tool_result": + inner, err := anthropicBlocks(b.Content, "tool_result.content") + if err != nil { + return nil, err + } + var result []string + for _, ib := range inner { + switch ib.Type { + case "text": + result = append(result, ib.Text) + case "image": + p, err := imagePart(ctx, ib.Source) + if err != nil { + return nil, err + } + parts = append(parts, p) + default: + return nil, invalid("tool_result: block type %q is not supported", ib.Type) + } + } + text := strings.Join(result, "\n") + if b.IsError { + text = "Error: " + text + } + out = append(out, &pb.ChatMessage{Role: pb.MessageRole_MESSAGE_ROLE_TOOL, Text: text, ToolCallId: b.ToolUseID}) + default: + return nil, invalid("block type %q is not supported in a user message", b.Type) + } + } + if len(parts) > 0 { + out = append(out, &pb.ChatMessage{Role: pb.MessageRole_MESSAGE_ROLE_USER, Text: strings.Join(texts, "\n\n"), Parts: parts}) + } + return out, nil +} + +func assistantMessage(blocks []anthropicBlock) ([]*pb.ChatMessage, error) { + msg := &pb.ChatMessage{Role: pb.MessageRole_MESSAGE_ROLE_ASSISTANT} + var texts []string + for _, b := range blocks { + switch { + case b.Type == "text": + texts = append(texts, b.Text) + case b.Type == "tool_use": + args := "{}" + if len(b.Input) > 0 && string(b.Input) != "null" { + args = string(b.Input) + } + msg.ToolCalls = append(msg.ToolCalls, &pb.ToolCall{ + Id: b.ID, Type: "function", Function: &pb.FunctionCall{Name: b.Name, Arguments: args}, + }) + case skippedBlocks[b.Type]: + default: + return nil, invalid("block type %q is not supported in an assistant message", b.Type) + } + } + msg.Text = strings.Join(texts, "\n\n") + if msg.Text == "" && len(msg.ToolCalls) == 0 { + return nil, nil + } + return []*pb.ChatMessage{msg}, nil +} + +// mediaBytes reads a base64 or https source; url sources go through safefetch +// like image_url on chat completions. +func mediaBytes(ctx context.Context, src *anthropicSource, what string) (mime string, data []byte, err error) { + if src == nil { + return "", nil, invalid("%s: source is required", what) + } + switch src.Type { + case "base64": + if data, err = base64.StdEncoding.DecodeString(src.Data); err != nil { + return "", nil, invalid("%s: invalid base64 data: %v", what, err) + } + mime = src.MediaType + case "url": + if !strings.HasPrefix(src.URL, "https://") { + return "", nil, invalid("%s: url %q must be https", what, src.URL) + } + if data, mime, err = safefetch.FetchURL(ctx, src.URL, maxRequestBytes); err != nil { + return "", nil, invalid("%s: failed to fetch %s: %v", what, src.URL, err) + } + default: + return "", nil, invalid("%s: source type %q is not supported, use base64 or url", what, src.Type) + } + if err := safefetch.ValidateMediaMimeType(mime); err != nil { + return "", nil, invalid("%s: %v", what, err) + } + return mime, data, nil +} + +func imagePart(ctx context.Context, src *anthropicSource) (*pb.ContentPart, error) { + mime, data, err := mediaBytes(ctx, src, "image") + if err != nil { + return nil, err + } + return &pb.ContentPart{Part: &pb.ContentPart_Image{Image: &pb.ImagePart{ + MimeType: mime, Data: data, Url: src.URL, Detail: "auto", + }}}, nil +} + +// documentPart returns a document part, or the text of a plain-text or +// content-block document. +func documentPart(ctx context.Context, b anthropicBlock) (*pb.ContentPart, string, error) { + if b.Source != nil { + switch b.Source.Type { + case "text": + return nil, b.Source.Data, nil + case "content": + inner, err := anthropicBlocks(b.Source.Content, "document.source.content") + if err != nil { + return nil, "", err + } + var texts []string + for _, ib := range inner { + if ib.Type != "text" { + return nil, "", invalid("document: content block type %q is not supported", ib.Type) + } + texts = append(texts, ib.Text) + } + return nil, strings.Join(texts, "\n\n"), nil + } + } + mime, data, err := mediaBytes(ctx, b.Source, "document") + if err != nil { + return nil, "", err + } + return &pb.ContentPart{Part: &pb.ContentPart_Document{Document: &pb.DocumentPart{ + MimeType: mime, Data: data, FileName: b.Title, + }}}, "", nil +} + +func (a *OpenAIAPI) messages(ctx context.Context, w http.ResponseWriter, r *http.Request) error { + var req messagesRequest + if err := decodeBody(r, &req); err != nil { + return err + } + if req.MaxTokens < 1 { + return invalid("max_tokens: must be at least 1") + } + preq, err := req.toProto(ctx, getPlacementHeader(r)) + if err != nil { + return err + } + if req.EnableCache != nil { + preq.EnableCache = *req.EnableCache + } else if h := getCacheHeader(r); h != "" { + preq.EnableCache = h == "true" || h == "1" + } + + id := "msg_" + randomID() + if req.Stream { + return a.streamMessages(ctx, w, preq, id) + } + resp, err := a.proxy.GenerateText(ctx, preq) + if err != nil { + return err + } + setRouteHeaders(w, resp) + setCacheHeader(w, preq.EnableCache, resp.Cached) + + var text, finish string + var calls []*pb.ToolCall + if len(resp.Details) > 0 { + d := resp.Details[0] + text, finish, calls = d.Text, d.FinishReason, d.ToolCalls + } else if len(resp.Choices) > 0 { + text = resp.Choices[0] + } + content := []map[string]any{} + if text != "" { + content = append(content, map[string]any{"type": "text", "text": text}) + } + for _, tc := range calls { + content = append(content, map[string]any{ + "type": "tool_use", "id": toolUseID(tc.GetId()), "name": tc.GetFunction().GetName(), + "input": toolInput(tc.GetFunction().GetArguments()), + }) + } + writeJSON(w, http.StatusOK, map[string]any{ + "id": id, "type": "message", "role": "assistant", "model": resp.ResolvedModel, + "content": content, "stop_reason": stopReason(finish, len(calls) > 0), "stop_sequence": nil, + "usage": anthropicUsage(resp.GetUsage(), 0), + }) + return nil +} + +// toolUseID keeps the vendor's call id, which the client sends back in its +// tool_result, and makes one up for vendors that give none. +func toolUseID(id string) string { + if id == "" { + return "toolu_" + randomID() + } + return id +} + +// toolInput is a tool call's arguments as the JSON object Messages clients +// expect; arguments that are not one become {}. +func toolInput(args string) json.RawMessage { + var obj map[string]json.RawMessage + if strings.TrimSpace(args) == "" || json.Unmarshal([]byte(args), &obj) != nil || obj == nil { + if strings.TrimSpace(args) != "" { + log.Warn().Int("bytes", len(args)).Msg("messages: tool call arguments are not a JSON object; sending {}") + } + return json.RawMessage("{}") + } + return json.RawMessage(args) +} + +func stopReason(finish string, toolCalls bool) string { + switch { + case finish == "length": + return "max_tokens" + case finish == "content_filter": + return "refusal" + case toolCalls || finish == "tool_calls": + return "tool_use" + default: + return "end_turn" + } +} + +// anthropicUsage splits the cached tokens out of the input count, as the +// Messages API reports them; estimate stands in before the vendor reports. +func anthropicUsage(u *pb.TokensUsage, estimate int64) map[string]any { + input := u.GetInput() - u.GetCacheRead() - u.GetCacheWrite() + if u == nil { + input = estimate + } + return map[string]any{ + "input_tokens": max(input, 0), + "output_tokens": u.GetOutput(), + "cache_read_input_tokens": u.GetCacheRead(), + "cache_creation_input_tokens": u.GetCacheWrite(), + } +} + +func (a *OpenAIAPI) streamMessages(ctx context.Context, w http.ResponseWriter, preq *pb.GenerateTextRequest, id string) error { + flusher, ok := w.(http.Flusher) + if !ok { + return status.Error(codes.Internal, "streaming unsupported by response writer") + } + s := &messagesStream{ + ctx: ctx, w: w, flusher: flusher, id: id, block: -1, + estimate: grpcapi.EstimateTokens(preq), enableCache: preq.EnableCache, toolBlocks: map[int32]bool{}, + } + err := a.proxy.GenerateTextStream(&pb.GenerateTextStreamRequest{Request: preq}, s) + if err != nil && !s.started { + return err + } + if err == nil { + s.start(nil) + s.finish() + return nil + } + // Headers are out: the error goes in the stream, as the Messages API sends it. + st, _ := status.FromError(err) + code, _ := httpStatus(st.Code()) + msg := st.Message() + if st.Code() == codes.Internal { + msg = "internal error" + } + s.event("error", map[string]any{"type": "error", "error": map[string]any{"type": anthropicErrorType(code), "message": msg}}) + return nil +} + +// messagesStream renders GenerateText chunks as Messages SSE events. Text +// and each tool call are content blocks, opened and closed in order. +type messagesStream struct { + grpc.ServerStream + ctx context.Context + w http.ResponseWriter + flusher http.Flusher + id string + estimate int64 + enableCache bool + cacheHit bool + started bool + model string + block int // index of the last block opened + open bool // whether that block is still open + blockTool bool + tool int32 // vendor index of the open tool call + toolBlocks map[int32]bool + toolCalls bool + finishedAs string + usage *pb.TokensUsage +} + +func (s *messagesStream) Context() context.Context { return s.ctx } + +func (s *messagesStream) event(name string, data any) { + b, _ := json.Marshal(data) + _, _ = fmt.Fprintf(s.w, "event: %s\ndata: %s\n\n", name, b) + s.flusher.Flush() +} + +func (s *messagesStream) start(chunk *pb.GenerateTextChunk) { + if s.started { + return + } + if chunk != nil { + setRouteHeaders(s.w, chunk) + s.model = chunk.ResolvedModel + } + setCacheHeader(s.w, s.enableCache, s.cacheHit) + startEventStream(s.w, s.flusher) + s.started = true + s.event("message_start", map[string]any{"type": "message_start", "message": map[string]any{ + "id": s.id, "type": "message", "role": "assistant", "model": s.model, "content": []any{}, + "stop_reason": nil, "stop_sequence": nil, "usage": anthropicUsage(nil, s.estimate), + }}) + s.event("ping", map[string]any{"type": "ping"}) +} + +func (s *messagesStream) openBlock(block map[string]any, tool bool, index int32) { + s.closeBlock() + s.block++ + s.open, s.blockTool, s.tool = true, tool, index + s.event("content_block_start", map[string]any{"type": "content_block_start", "index": s.block, "content_block": block}) +} + +func (s *messagesStream) closeBlock() { + if !s.open { + return + } + s.event("content_block_stop", map[string]any{"type": "content_block_stop", "index": s.block}) + s.open = false +} + +func (s *messagesStream) delta(delta map[string]any) { + s.event("content_block_delta", map[string]any{"type": "content_block_delta", "index": s.block, "delta": delta}) +} + +func (s *messagesStream) Send(chunk *pb.GenerateTextChunk) error { + if chunk.Cached { + s.cacheHit = true + } + s.start(chunk) + if chunk.Delta != "" { + if !s.open || s.blockTool { + s.openBlock(map[string]any{"type": "text", "text": ""}, false, -1) + } + s.delta(map[string]any{"type": "text_delta", "text": chunk.Delta}) + } + for _, tc := range chunk.ToolCallChunks { + if !s.open || !s.blockTool || s.tool != tc.Index { + if s.toolBlocks[tc.Index] { + return status.Errorf(codes.Internal, "tool call %d continued after the next content block began", tc.Index) + } + s.toolBlocks[tc.Index] = true + s.toolCalls = true + s.openBlock(map[string]any{"type": "tool_use", "id": toolUseID(tc.Id), "name": tc.Name, "input": map[string]any{}}, true, tc.Index) + } + if tc.ArgumentsDelta != "" { + s.delta(map[string]any{"type": "input_json_delta", "partial_json": tc.ArgumentsDelta}) + } + } + if chunk.FinishReason != "" { + s.finishedAs = chunk.FinishReason + } + if chunk.Usage != nil { + s.usage = chunk.Usage + } + return nil +} + +func (s *messagesStream) finish() { + s.closeBlock() + usage := anthropicUsage(s.usage, s.estimate) + s.event("message_delta", map[string]any{"type": "message_delta", + "delta": map[string]any{"stop_reason": stopReason(s.finishedAs, s.toolCalls), "stop_sequence": nil}, + "usage": usage}) + s.event("message_stop", map[string]any{"type": "message_stop"}) +} + +func writeAnthropicError(w http.ResponseWriter, err error) { + code, _, msg, _ := apiError(w, err) + writeJSON(w, code, map[string]any{"type": "error", "error": map[string]any{ + "type": anthropicErrorType(code), "message": msg, + }}) +} + +// anthropicErrorType names the Messages API error for an HTTP status; the +// statuses themselves are those of the HTTP error table. +func anthropicErrorType(code int) string { + switch code { + case http.StatusBadRequest: + return "invalid_request_error" + case http.StatusUnauthorized: + return "authentication_error" + case http.StatusForbidden: + return "permission_error" + case http.StatusNotFound: + return "not_found_error" + case http.StatusRequestEntityTooLarge: + return "request_too_large" + case http.StatusTooManyRequests: + return "rate_limit_error" + case http.StatusServiceUnavailable: + return "overloaded_error" + case http.StatusGatewayTimeout: + return "timeout_error" + default: + return "api_error" + } +} + +// countTokens estimates like the token budgets do (a rune count), adding the +// tool definitions and tool calls the estimate leaves out. +func (a *OpenAIAPI) countTokens(ctx context.Context, w http.ResponseWriter, r *http.Request) error { + var req messagesRequest + if err := decodeBody(r, &req); err != nil { + return err + } + preq, err := req.toProto(ctx, "") + if err != nil { + return err + } + n := grpcapi.EstimateTokens(preq) + extra := 0 + for _, t := range req.Tools { + extra += utf8.RuneCountInString(t.Name) + utf8.RuneCountInString(t.Description) + utf8.RuneCount(t.InputSchema) + } + for _, m := range preq.Messages { + for _, tc := range m.ToolCalls { + extra += utf8.RuneCountInString(tc.GetFunction().GetName()) + utf8.RuneCountInString(tc.GetFunction().GetArguments()) + } + } + setObservabilityHeader(w, "Token-Count", "estimate") + writeJSON(w, http.StatusOK, map[string]any{"input_tokens": n + int64(extra/4)}) + return nil +} diff --git a/internal/server/anthropic_test.go b/internal/server/anthropic_test.go new file mode 100644 index 0000000..cfcbb9a --- /dev/null +++ b/internal/server/anthropic_test.go @@ -0,0 +1,285 @@ +package server + +import ( + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + pb "github.com/ai-process/ktms/gen/ktms/v1" +) + +// messagesCall sends the key the way the Anthropic SDKs do, in x-api-key. +func messagesCall(h http.Handler, path, key, body string) *httptest.ResponseRecorder { + req := httptest.NewRequest("POST", path, strings.NewReader(body)) + if key != "" { + req.Header.Set("X-Api-Key", key) + } + req.Header.Set("Anthropic-Version", "2023-06-01") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec +} + +func anthropicErr(t *testing.T, rec *httptest.ResponseRecorder) (typ, msg string) { + t.Helper() + var out struct { + Type string `json:"type"` + Error struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || out.Type != "error" { + t.Fatalf("not an Anthropic error body: %s", rec.Body) + } + return out.Error.Type, out.Error.Message +} + +func TestMessagesMapsRequest(t *testing.T) { + fp, h, key := newTestAPI(t) + png := base64.StdEncoding.EncodeToString([]byte("\x89PNG\r\n\x1a\n")) + rec := messagesCall(h, "/v1/messages?beta=true", key, `{ + "model":"auto:high","max_tokens":32000, + "system":[{"type":"text","text":"You are Claude Code."},{"type":"text","text":"Be brief.","cache_control":{"type":"ephemeral"}}], + "metadata":{"user_id":"user_abc_session_1"}, + "temperature":1,"thinking":{"type":"enabled","budget_tokens":1024}, + "tools":[{"name":"Read","description":"Read a file","input_schema":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}], + "tool_choice":{"type":"any"}, + "messages":[ + {"role":"user","content":"read main.go"}, + {"role":"assistant","content":[ + {"type":"thinking","thinking":"hm","signature":"sig"}, + {"type":"text","text":"Reading."}, + {"type":"tool_use","id":"toolu_1","name":"Read","input":{"path":"main.go"}}]}, + {"role":"user","content":[ + {"type":"tool_result","tool_use_id":"toolu_1","content":[{"type":"text","text":"package main"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"`+png+`"}}]}, + {"type":"text","text":"x"}, + {"type":"text","text":"explain it"}]}]}`) + if rec.Code != 200 { + t.Fatalf("code %d: %s", rec.Code, rec.Body) + } + r := fp.text + if r.Effort != pb.Effort_EFFORT_HIGH || !r.AutoRoute || r.MaxOutputTokens != 32000 { + t.Errorf("routing: %+v", r) + } + if r.BaseSystemInstruction != "You are Claude Code.\n\nBe brief." || r.Attributes["user_id"] != "user_abc_session_1" { + t.Errorf("system %q attrs %v", r.BaseSystemInstruction, r.Attributes) + } + if len(r.Tools) != 1 || r.Tools[0].Function.Name != "Read" || r.Tools[0].Function.Parameters.Properties["path"].Type != "string" || + r.ToolChoice.GetMode() != pb.ToolChoice_MODE_REQUIRED { + t.Errorf("tools %v choice %v", r.Tools, r.ToolChoice) + } + if len(r.Messages) != 4 { + t.Fatalf("messages: %v", r.Messages) + } + asst, tool, user := r.Messages[1], r.Messages[2], r.Messages[3] + if asst.Text != "Reading." || len(asst.ToolCalls) != 1 || asst.ToolCalls[0].Id != "toolu_1" || + asst.ToolCalls[0].Function.Arguments != `{"path":"main.go"}` { + t.Errorf("assistant: %v", asst) + } + if tool.Role != pb.MessageRole_MESSAGE_ROLE_TOOL || tool.ToolCallId != "toolu_1" || tool.Text != "package main" { + t.Errorf("tool result: %v", tool) + } + // The tool result's image rides in the user message after it. + if user.Role != pb.MessageRole_MESSAGE_ROLE_USER || user.Text != "x\n\nexplain it" || + len(user.Parts) != 3 || user.Parts[0].GetImage().GetMimeType() != "image/png" { + t.Errorf("user: %v", user) + } +} + +func TestMessagesResponse(t *testing.T) { + fp, h, key := newTestAPI(t) + fp.respDetails = []*pb.Choice{{ + Text: "Let me look.", + ToolCalls: []*pb.ToolCall{ + {Id: "call_9", Function: &pb.FunctionCall{Name: "Read", Arguments: `{"path":"a.go"}`}}, + {Function: &pb.FunctionCall{Name: "Glob", Arguments: ""}}, + }, + FinishReason: "tool_calls", + }} + rec := messagesCall(h, "/v1/messages", key, `{"model":"m1","max_tokens":100,"messages":[{"role":"user","content":"hi"}]}`) + if rec.Code != 200 { + t.Fatalf("code %d: %s", rec.Code, rec.Body) + } + var out struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Model string `json:"model"` + StopReason string `json:"stop_reason"` + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + ID string `json:"id"` + Name string `json:"name"` + Input json.RawMessage `json:"input"` + } `json:"content"` + Usage map[string]int `json:"usage"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatal(err) + } + if !strings.HasPrefix(out.ID, "msg_") || out.Type != "message" || out.Role != "assistant" || out.Model != "m1" || + out.StopReason != "tool_use" || len(out.Content) != 3 { + t.Fatalf("message: %s", rec.Body) + } + if out.Content[0].Text != "Let me look." || out.Content[1].ID != "call_9" || string(out.Content[1].Input) != `{"path":"a.go"}` || + !strings.HasPrefix(out.Content[2].ID, "toolu_") || string(out.Content[2].Input) != "{}" { + t.Errorf("content: %s", rec.Body) + } + if out.Usage["input_tokens"] != 3 || out.Usage["output_tokens"] != 2 { + t.Errorf("usage: %v", out.Usage) + } + if rec.Header().Get("X-KTMS-Vendor") != "openai" || rec.Header().Get("X-KTMS-Rule") != "default" { + t.Errorf("headers: %v", rec.Header()) + } +} + +func TestMessagesUsageSplitsCache(t *testing.T) { + u := anthropicUsage(&pb.TokensUsage{Input: 100, Output: 7, CacheRead: 60, CacheWrite: 30}, 0) + if u["input_tokens"] != int64(10) || u["cache_read_input_tokens"] != int64(60) || u["cache_creation_input_tokens"] != int64(30) { + t.Errorf("usage %v", u) + } + if u := anthropicUsage(nil, 42); u["input_tokens"] != int64(42) || u["output_tokens"] != int64(0) { + t.Errorf("estimate %v", u) + } +} + +// sseEvents returns the event names of a stream, in order. +func sseEvents(body string) []string { + var names []string + for _, line := range strings.Split(body, "\n") { + if name, ok := strings.CutPrefix(line, "event: "); ok { + names = append(names, name) + } + } + return names +} + +func TestMessagesStream(t *testing.T) { + fp, h, key := newTestAPI(t) + fp.streamChunks = []*pb.GenerateTextChunk{ + {Delta: "Checking", ResolvedModel: "m1", ResolvedVendor: "openai", MatchedRule: "default"}, + {Delta: " now."}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, Id: "call_a", Name: "Read", ArgumentsDelta: `{"path":`}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, ArgumentsDelta: `"a.go"}`}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 1, Name: "Glob", ArgumentsDelta: `{}`}}}, + {FinishReason: "tool_calls", Usage: &pb.TokensUsage{Input: 50, Output: 9, CacheRead: 40}}, + } + rec := messagesCall(h, "/v1/messages", key, `{"model":"auto","max_tokens":100,"stream":true,"messages":[{"role":"user","content":"x"}]}`) + if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" || rec.Header().Get("X-KTMS-Vendor") != "openai" { + t.Fatalf("headers %v: %s", rec.Header(), rec.Body) + } + b := rec.Body.String() + want := []string{"message_start", "ping", + "content_block_start", "content_block_delta", "content_block_delta", "content_block_stop", + "content_block_start", "content_block_delta", "content_block_delta", "content_block_stop", + "content_block_start", "content_block_delta", "content_block_stop", + "message_delta", "message_stop"} + if got := sseEvents(b); strings.Join(got, ",") != strings.Join(want, ",") { + t.Fatalf("events %v\n%s", got, b) + } + for _, s := range []string{ + `"model":"m1"`, + `{"content_block":{"text":"","type":"text"},"index":0,"type":"content_block_start"}`, + `"delta":{"text":" now.","type":"text_delta"}`, + `"content_block":{"id":"call_a","input":{},"name":"Read","type":"tool_use"},"index":1`, + `"delta":{"partial_json":"\"a.go\"}","type":"input_json_delta"},"index":1`, + `"name":"Glob"`, + `"stop_reason":"tool_use"`, + `"cache_read_input_tokens":40,"input_tokens":10,"output_tokens":9`, + } { + if !strings.Contains(b, s) { + t.Errorf("stream missing %s:\n%s", s, b) + } + } +} + +func TestMessagesStreamToolOutOfOrder(t *testing.T) { + fp, h, key := newTestAPI(t) + fp.streamChunks = []*pb.GenerateTextChunk{ + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, Id: "a", Name: "Read"}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 1, Id: "b", Name: "Glob"}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, ArgumentsDelta: "{}"}}}, + } + rec := messagesCall(h, "/v1/messages", key, `{"max_tokens":10,"stream":true,"messages":[{"role":"user","content":"x"}]}`) + if ev := sseEvents(rec.Body.String()); ev[len(ev)-1] != "error" { + t.Errorf("events %v", ev) + } +} + +func TestMessagesStreamErrors(t *testing.T) { + fp, h, key := newTestAPI(t) + fp.textErr = status.Error(codes.ResourceExhausted, "throttled:rpm") + rec := messagesCall(h, "/v1/messages", key, `{"max_tokens":10,"stream":true,"messages":[{"role":"user","content":"x"}]}`) + if typ, _ := anthropicErr(t, rec); rec.Code != 429 || typ != "rate_limit_error" || rec.Header().Get("Retry-After") == "" { + t.Errorf("before the stream: %d %s %v", rec.Code, typ, rec.Header()) + } + + fp.textErr = nil + fp.emptyStream = true + rec = messagesCall(h, "/v1/messages", key, `{"max_tokens":10,"stream":true,"messages":[{"role":"user","content":"x"}]}`) + if got := strings.Join(sseEvents(rec.Body.String()), ","); got != "message_start,ping,message_delta,message_stop" { + t.Errorf("no chunks: %s", got) + } +} + +func TestMessagesErrors(t *testing.T) { + fp, h, key := newTestAPI(t) + cases := []struct { + key, body, typ string + code int + }{ + {"", `{"max_tokens":1,"messages":[{"role":"user","content":"x"}]}`, "authentication_error", 401}, + {key, `{"messages":[{"role":"user","content":"x"}]}`, "invalid_request_error", 400}, + {key, `{"max_tokens":1,"messages":[]}`, "invalid_request_error", 400}, + {key, `{"max_tokens":1,"messages":[{"role":"system","content":"x"}]}`, "invalid_request_error", 400}, + {key, `{"max_tokens":1,"tools":[{"type":"web_search_20250305","name":"web_search"}],"messages":[{"role":"user","content":"x"}]}`, "invalid_request_error", 400}, + {key, `{"max_tokens":1,"messages":[{"role":"user","content":[{"type":"image","source":{"type":"url","url":"http://169.254.169.254/x.png"}}]}]}`, "invalid_request_error", 400}, + {key, `{"max_tokens":1,"messages":[{"role":"user","content":[{"type":"image","source":{"type":"file","file_id":"f"}}]}]}`, "invalid_request_error", 400}, + } + for _, c := range cases { + rec := messagesCall(h, "/v1/messages", c.key, c.body) + if typ, msg := anthropicErr(t, rec); rec.Code != c.code || typ != c.typ { + t.Errorf("%s: %d %s %q", c.body, rec.Code, typ, msg) + } + } + + // Authorization works as well as x-api-key, as ANTHROPIC_AUTH_TOKEN sends it. + if rec := do(h, "POST", "/v1/messages", key, `{"max_tokens":1,"messages":[{"role":"user","content":"x"}]}`); rec.Code != 200 { + t.Errorf("bearer: %d %s", rec.Code, rec.Body) + } + // x-api-key is for the Messages API only. + if rec := messagesCall(h, "/v1/chat/completions", key, `{"messages":[{"role":"user","content":"x"}]}`); rec.Code != 401 { + t.Errorf("x-api-key on chat completions: %d", rec.Code) + } + + fp.textErr = status.Error(codes.Unavailable, "vendors_unavailable") + rec := messagesCall(h, "/v1/messages", key, `{"max_tokens":1,"messages":[{"role":"user","content":"x"}]}`) + if typ, msg := anthropicErr(t, rec); rec.Code != 503 || typ != "overloaded_error" || msg != "vendors_unavailable" { + t.Errorf("unavailable: %d %s %s", rec.Code, typ, msg) + } +} + +func TestCountTokens(t *testing.T) { + _, h, key := newTestAPI(t) + rec := messagesCall(h, "/v1/messages/count_tokens", key, `{"model":"auto", + "system":"`+strings.Repeat("s", 400)+`", + "tools":[{"name":"Read","input_schema":{"type":"object"}}], + "messages":[{"role":"user","content":"`+strings.Repeat("u", 400)+`"}]}`) + var out struct { + InputTokens int64 `json:"input_tokens"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || rec.Code != 200 { + t.Fatalf("%d %s", rec.Code, rec.Body) + } + if out.InputTokens < 200 || out.InputTokens > 220 || rec.Header().Get("X-KTMS-Token-Count") != "estimate" { + t.Errorf("count %d headers %v", out.InputTokens, rec.Header()) + } +} diff --git a/internal/server/openai.go b/internal/server/openai.go index ae9a477..0217779 100644 --- a/internal/server/openai.go +++ b/internal/server/openai.go @@ -66,6 +66,8 @@ func (a *OpenAIAPI) Register(mux *http.ServeMux) { mux.HandleFunc("POST /v1/decisions/rollouts", a.wrap(judge, a.decisionsRollouts)) mux.HandleFunc("GET /v1/models", a.wrap(models, a.listModels)) mux.HandleFunc("GET /v1/models/{id...}", a.wrap(models, a.getModel)) + mux.HandleFunc("POST /v1/messages", a.serve(chat, a.messages, writeAnthropicError, true)) + mux.HandleFunc("POST /v1/messages/count_tokens", a.serve(chat, a.countTokens, writeAnthropicError, true)) } type apiHandler func(ctx context.Context, w http.ResponseWriter, r *http.Request) error @@ -73,6 +75,12 @@ type apiHandler func(ctx context.Context, w http.ResponseWriter, r *http.Request // wrap authenticates, recovers panics and turns returned errors into // OpenAI-style error bodies. func (a *OpenAIAPI) wrap(fullMethod string, h apiHandler) http.HandlerFunc { + return a.serve(fullMethod, h, writeAPIError, false) +} + +// serve is wrap with the error body format and, for Anthropic clients, the +// key in x-api-key when there is no Authorization header. +func (a *OpenAIAPI) serve(fullMethod string, h apiHandler, writeErr func(http.ResponseWriter, error), apiKeyHeader bool) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { start := time.Now() var err error @@ -80,7 +88,7 @@ func (a *OpenAIAPI) wrap(fullMethod string, h apiHandler) http.HandlerFunc { if rec := recover(); rec != nil { panicsafe.Report("http handler "+fullMethod, rec) err = status.Error(codes.Internal, "internal error") - writeAPIError(w, err) + writeErr(w, err) } rpcLogEvent(err).Str("method", fullMethod).Str("transport", "http"). Dur("duration", time.Since(start)).Msg("http request") @@ -89,16 +97,18 @@ func (a *OpenAIAPI) wrap(fullMethod string, h apiHandler) http.HandlerFunc { md := metadata.MD{} if auth := r.Header.Get("Authorization"); auth != "" { md.Set(metadataKeyAuthorization, auth) + } else if key := r.Header.Get("X-Api-Key"); apiKeyHeader && key != "" { + md.Set(metadataKeyAuthorization, "Bearer "+key) } ctx := metadata.NewIncomingContext(r.Context(), md) id, aerr := authorize(ctx, a.verifier, fullMethod) if aerr != nil { err = aerr - writeAPIError(w, err) + writeErr(w, err) return } if err = h(apikeys.WithIdentity(ctx, id), w, r); err != nil { - writeAPIError(w, err) + writeErr(w, err) } } } @@ -698,26 +708,8 @@ func (a *OpenAIAPI) chatCompletions(ctx context.Context, w http.ResponseWriter, if err != nil { return err } - setResolvedHeaders(w, resp.ResolvedVendor, resp.MatchedRule) - if resp.Placement != "" { - setObservabilityHeader(w, "Placement", resp.Placement) - } - if resp.RouteMode != "" { - setObservabilityHeader(w, "Route-Mode", resp.RouteMode) - } - if resp.JudgeScore > 0 { - setObservabilityHeader(w, "Judge-Score", fmt.Sprintf("%.2f", resp.JudgeScore)) - } - if resp.JudgeModel != "" { - setObservabilityHeader(w, "Judge-Model", resp.JudgeModel) - } - if preq.EnableCache { - if resp.Cached { - setObservabilityHeader(w, "Cache", "HIT") - } else { - setObservabilityHeader(w, "Cache", "MISS") - } - } + setRouteHeaders(w, resp) + setCacheHeader(w, preq.EnableCache, resp.Cached) usage := usageJSON(resp.GetUsage()) choices := make([]map[string]any, 0, len(resp.Choices)) @@ -848,36 +840,22 @@ func (s *httpStreamAdapter) ensureHeaders(chunk *pb.GenerateTextChunk) { return } if chunk != nil { - setResolvedHeaders(s.w, chunk.ResolvedVendor, chunk.MatchedRule) - if chunk.Placement != "" { - setObservabilityHeader(s.w, "Placement", chunk.Placement) - } - if chunk.RouteMode != "" { - setObservabilityHeader(s.w, "Route-Mode", chunk.RouteMode) - } - if chunk.JudgeScore > 0 { - setObservabilityHeader(s.w, "Judge-Score", fmt.Sprintf("%.2f", chunk.JudgeScore)) - } - if chunk.JudgeModel != "" { - setObservabilityHeader(s.w, "Judge-Model", chunk.JudgeModel) - } - } - if s.enableCache { - if s.cacheHit { - setObservabilityHeader(s.w, "Cache", "HIT") - } else { - setObservabilityHeader(s.w, "Cache", "MISS") - } + setRouteHeaders(s.w, chunk) } - s.w.Header().Set("Content-Type", "text/event-stream") - s.w.Header().Set("Cache-Control", "no-cache") - s.w.Header().Set("Connection", "keep-alive") - s.w.Header().Set("X-Accel-Buffering", "no") - s.w.WriteHeader(http.StatusOK) - s.flusher.Flush() + setCacheHeader(s.w, s.enableCache, s.cacheHit) + startEventStream(s.w, s.flusher) s.headersSent = true } +func startEventStream(w http.ResponseWriter, flusher http.Flusher) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + flusher.Flush() +} + func (s *httpStreamAdapter) Send(chunk *pb.GenerateTextChunk) error { if chunk.Cached { s.cacheHit = true @@ -978,6 +956,36 @@ func setObservabilityHeader(w http.ResponseWriter, suffix, value string) { w.Header().Set(headerWarning, deprecationWarn) } +// routeInfo is how a text response or stream chunk says where it was routed. +type routeInfo interface { + GetResolvedVendor() string + GetMatchedRule() string + GetPlacement() string + GetRouteMode() string + GetJudgeScore() float64 + GetJudgeModel() string +} + +func setRouteHeaders(w http.ResponseWriter, r routeInfo) { + setResolvedHeaders(w, r.GetResolvedVendor(), r.GetMatchedRule()) + setObservabilityHeader(w, "Placement", r.GetPlacement()) + setObservabilityHeader(w, "Route-Mode", r.GetRouteMode()) + if r.GetJudgeScore() > 0 { + setObservabilityHeader(w, "Judge-Score", fmt.Sprintf("%.2f", r.GetJudgeScore())) + } + setObservabilityHeader(w, "Judge-Model", r.GetJudgeModel()) +} + +func setCacheHeader(w http.ResponseWriter, enabled, hit bool) { + switch { + case !enabled: + case hit: + setObservabilityHeader(w, "Cache", "HIT") + default: + setObservabilityHeader(w, "Cache", "MISS") + } +} + func setResolvedHeaders(w http.ResponseWriter, vendor, rule string) { if vendor != "" { setObservabilityHeader(w, "Vendor", vendor) @@ -1953,11 +1961,19 @@ func httpStatus(c codes.Code) (int, string) { } func writeAPIError(w http.ResponseWriter, err error) { + code, typ, msg, errCode := apiError(w, err) + writeJSON(w, code, map[string]any{"error": map[string]any{ + "message": msg, "type": typ, "param": nil, "code": errCode, + }}) +} + +// apiError resolves err into the HTTP status, OpenAI error type, message and +// machine-readable code, setting Retry-After when the client should wait. +func apiError(w http.ResponseWriter, err error) (code int, typ, msg string, errCode any) { st, _ := status.FromError(err) - code, typ := httpStatus(st.Code()) - msg := st.Message() + code, typ = httpStatus(st.Code()) + msg = st.Message() // Machine-readable reasons ("throttled:rpm", "no_capable_model") become the code. - var errCode any if reasonRe.MatchString(msg) { errCode = msg } else if strings.HasPrefix(msg, "guardrail_violation") { @@ -1985,9 +2001,7 @@ func writeAPIError(w http.ResponseWriter, err error) { if msg == "vendor_error" { code = http.StatusBadGateway } - writeJSON(w, code, map[string]any{"error": map[string]any{ - "message": msg, "type": typ, "param": nil, "code": errCode, - }}) + return code, typ, msg, errCode } // retryAfterSeconds reads a RetryInfo on the status: spend limits wait for diff --git a/plugin/skills/integrate/SKILL.md b/plugin/skills/integrate/SKILL.md index fa85950..79877d4 100644 --- a/plugin/skills/integrate/SKILL.md +++ b/plugin/skills/integrate/SKILL.md @@ -5,7 +5,8 @@ description: Move a service's LLM calls onto Keystone (KTMS), the self-hosted LL # Connect a service to Keystone -Keystone speaks the OpenAI wire format on HTTP `/v1` and its own gRPC API on `:9090`. A service +Keystone speaks the OpenAI wire format on HTTP `/v1`, the Anthropic Messages API on +`/v1/messages`, and its own gRPC API on `:9090`. A service needs a client key (`ktms_…`, scope `generate`) and the gateway's address; it never names vendor models, it asks for an effort and Keystone's rules pick the chain. @@ -24,6 +25,8 @@ models, it asks for an effort and Keystone's rules pick the chain. - One small client module owns the base URL, key, timeouts, retries and attribution. Read `KTMS_BASE_URL` (or `KTMS_GRPC_ADDR`) and `KTMS_API_KEY` from the environment; add both to the example env file and the deployment config. +- A service on the Anthropic SDK can keep it: its base URL is Keystone without `/v1`, its key + the KTMS key, and `metadata.user_id` is the attribution. - Model values: `auto` (a judge picks the effort), `auto:low|medium|high`, or `auto::` with placement `nodes-first`, `nodes-only`, `cloud-first`, `cloud-only`. In gRPC set `effort`. An exact model id bypasses the rules; embeddings always need @@ -49,4 +52,5 @@ models, it asks for an effort and Keystone's rules pick the chain. - `ktms chat --effort low "hello"` with the service's key shows the rule and vendor that served it. Sources: [OpenAI-compatible HTTP API](https://github.com/ai-process/ktms/blob/main/README.md#openai-compatible-http-api), +[Anthropic Messages API](https://github.com/ai-process/ktms/blob/main/README.md#anthropic-messages-api), [Calling Keystone](https://github.com/ai-process/ktms/blob/main/README.md#calling-keystone), [Error contract](https://github.com/ai-process/ktms/blob/main/README.md#error-contract). From 8095e435f40456ea7286224f27bfb40e34175613 Mon Sep 17 00:00:00 2001 From: Vasil Date: Wed, 28 Oct 2026 01:27:17 +0100 Subject: [PATCH 2/2] fix(http): review findings on /v1/messages Streamed tool calls open once their name is complete and send no null arguments, stream write errors stop the vendor stream, mid-stream errors use the shared error mapping, count_tokens does not fetch url sources, Claude Code's session-scoped user_id keeps its device or account as the user, and budget estimates count tools and tool calls. --- CHANGELOG.md | 6 +- README.md | 8 +- internal/grpcapi/handler.go | 8 + internal/grpcapi/handler_test.go | 8 + internal/server/anthropic.go | 241 ++++++++++++++++++++---------- internal/server/anthropic_test.go | 68 ++++++++- internal/server/openai.go | 17 ++- internal/server/openai_test.go | 3 +- 8 files changed, 262 insertions(+), 97 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 80ad6cc..0e4d9b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,11 @@ All notable changes to KTMS are listed here. The format follows fallback, limits, guardrails, cache, `X-KTMS-*` headers and usage events with chat completions, answers in the Messages format whichever vendor's model serves it, streams the Messages SSE events, and returns errors in the Messages shape. `POST /v1/messages/count_tokens` returns an - estimate. + estimate. Claude Code's session-scoped `metadata.user_id` becomes its account or device as the + user, with the session in `attributes["session_id"]`. + +### Changed +- Token budget reservations count tool definitions and the tool calls in the conversation. ## [1.4.0] - 2026-10-27 diff --git a/README.md b/README.md index 81526c3..c3cd36e 100644 --- a/README.md +++ b/README.md @@ -1441,7 +1441,7 @@ print(reply.content[0].text, reply.model) # `model` is the one that answered | `tools`, `tool_choice` | Client tools with `input_schema` (the same JSON Schema subset as `response_format`); `tool_choice` `auto`, `any`, `tool` and `none`. Server tools such as `web_search` are rejected with `400`. | | `max_tokens` | Required, the output ceiling. | | `stream` | SSE events `message_start`, `ping`, `content_block_start`, `content_block_delta` (`text_delta`, `input_json_delta`), `content_block_stop`, `message_delta` and `message_stop`. `message_start` carries an input token estimate; `message_delta` carries the vendor's counts. | -| `metadata.user_id` | `attributes["user_id"]`, as `user` on chat completions; other `metadata` keys are attributes, and `action_id` is the usage action label. | +| `metadata.user_id` | `attributes["user_id"]`, as `user` on chat completions; other `metadata` keys are attributes, and `action_id` is the usage action label. Claude Code's `user_id` names a session: Keystone keeps its account (or, without one, its device) as the user and puts the session in `attributes["session_id"]`, so per-user limits and reports span sessions. | | `enable_cache`, `placement` | Extensions, as on chat completions; so are the `X-KTMS-Cache` and `X-KTMS-Placement` request headers (set them in Claude Code with `ANTHROPIC_CUSTOM_HEADERS`). | | `temperature`, `top_p`, `top_k`, `stop_sequences`, `thinking`, `service_tier` | Accepted and ignored. | @@ -1450,8 +1450,10 @@ The key goes in `x-api-key` or `Authorization: Bearer`. The response's `stop_rea `cache_creation_input_tokens` out of `input_tokens`. Tool call ids are the vendor's, or `toolu_…` when the vendor gives none. -`POST /v1/messages/count_tokens` answers with an estimate, not a tokenizer's count: about four -characters per token, tool definitions included, with the `X-KTMS-Token-Count: estimate` header. +`POST /v1/messages/count_tokens` answers with the estimate the token budgets reserve, not a +tokenizer's count: about four characters per token, tool definitions included, a flat allowance +per image or document (`url` sources are not fetched), and the `X-KTMS-Token-Count: estimate` +header. Errors use the Messages shape, `{"type": "error", "error": {"type", "message"}}`, with the HTTP statuses of [the table above](#errors): `invalid_request_error` (400), `authentication_error` diff --git a/internal/grpcapi/handler.go b/internal/grpcapi/handler.go index 4e135d0..2bea4a4 100644 --- a/internal/grpcapi/handler.go +++ b/internal/grpcapi/handler.go @@ -8,6 +8,7 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" pb "github.com/ai-process/ktms/gen/ktms/v1" "github.com/ai-process/ktms/internal/apikeys" @@ -664,12 +665,19 @@ func buildChat(req *pb.GenerateTextRequest, id *apikeys.Identity, ruleName, effo // EstimateTokens is the budget reservation guess: total runes / 4 plus heuristic // allowances for media parts. Settled against vendor-reported actuals after the call. +// Tool definitions count by their encoded size: coding agents send thousands of tokens of them. func EstimateTokens(req *pb.GenerateTextRequest) int64 { runes := utf8.RuneCountInString(req.GetBaseSystemInstruction()) + utf8.RuneCountInString(req.GetRequestInstruction()) + for _, t := range req.GetTools() { + runes += proto.Size(t) + } var extraTokens int64 for _, m := range req.GetMessages() { runes += utf8.RuneCountInString(m.GetText()) + for _, tc := range m.GetToolCalls() { + runes += utf8.RuneCountInString(tc.GetFunction().GetName()) + utf8.RuneCountInString(tc.GetFunction().GetArguments()) + } for _, p := range m.GetParts() { switch part := p.GetPart().(type) { case *pb.ContentPart_Text: diff --git a/internal/grpcapi/handler_test.go b/internal/grpcapi/handler_test.go index 20671a2..5be73e4 100644 --- a/internal/grpcapi/handler_test.go +++ b/internal/grpcapi/handler_test.go @@ -20,6 +20,7 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" pb "github.com/ai-process/ktms/gen/ktms/v1" "github.com/ai-process/ktms/internal/apikeys" @@ -551,6 +552,13 @@ func TestEstimateTokens(t *testing.T) { if got := EstimateTokens(req); got != 4 { // 12/4 + 1 t.Fatalf("estimate = %d, want 4", got) } + req.Messages = append(req.Messages, &pb.ChatMessage{ToolCalls: []*pb.ToolCall{ + {Function: &pb.FunctionCall{Name: "Read", Arguments: `{"path":"a.go"}`}}, // 4 + 15 runes + }}) + req.Tools = []*pb.Tool{{Type: "function", Function: &pb.FunctionDefinition{Name: "Read", Description: strings.Repeat("d", 80)}}} + if got, want := EstimateTokens(req), int64(12+19+proto.Size(req.Tools[0]))/4+1; got != want { + t.Fatalf("estimate with tools = %d, want %d", got, want) + } } func TestGenerateTextAllThrottled(t *testing.T) { diff --git a/internal/server/anthropic.go b/internal/server/anthropic.go index 437ec8a..08aa690 100644 --- a/internal/server/anthropic.go +++ b/internal/server/anthropic.go @@ -7,7 +7,6 @@ import ( "fmt" "net/http" "strings" - "unicode/utf8" "github.com/rs/zerolog/log" "google.golang.org/grpc" @@ -98,7 +97,9 @@ var skippedBlocks = map[string]bool{ "thinking": true, "redacted_thinking": true, "server_tool_use": true, "web_search_tool_result": true, } -func (req *messagesRequest) toProto(ctx context.Context, headerPlacement string) (*pb.GenerateTextRequest, error) { +// toProto converts the request. Without fetch, url sources are not downloaded +// and stand in as empty parts, which is all a token estimate needs. +func (req *messagesRequest) toProto(ctx context.Context, headerPlacement string, fetch bool) (*pb.GenerateTextRequest, error) { if len(req.Messages) == 0 { return nil, invalid("messages: at least one message is required") } @@ -119,6 +120,12 @@ func (req *messagesRequest) toProto(ctx context.Context, headerPlacement string) EffortSpecified: dec.EffortSpecified, } out.Attributes, out.ActionId = attributes(req.Metadata, "") + if user, session := claudeCodeUser(out.Attributes["user_id"]); user != "" { + out.Attributes["user_id"] = user + if out.Attributes["session_id"] == "" { + out.Attributes["session_id"] = session + } + } system, err := anthropicBlocks(req.System, "system") if err != nil { @@ -174,7 +181,7 @@ func (req *messagesRequest) toProto(ctx context.Context, headerPlacement string) var msgs []*pb.ChatMessage switch m.Role { case "user": - msgs, err = userMessages(ctx, blocks) + msgs, err = userMessages(ctx, blocks, fetch) case "assistant": msgs, err = assistantMessage(blocks) default: @@ -188,6 +195,30 @@ func (req *messagesRequest) toProto(ctx context.Context, headerPlacement string) return out, nil } +// claudeCodeUser reads the metadata.user_id Claude Code sends, a JSON object +// or "user__account__session_", whose session part would +// otherwise make every session its own user for limits and reports. +func claudeCodeUser(id string) (user, session string) { + if strings.HasPrefix(id, "{") { + var cc struct { + DeviceID string `json:"device_id"` + AccountUUID string `json:"account_uuid"` + SessionID string `json:"session_id"` + } + if json.Unmarshal([]byte(id), &cc) != nil { + return "", "" + } + if cc.AccountUUID != "" { + return cc.AccountUUID, cc.SessionID + } + return cc.DeviceID, cc.SessionID + } + if before, after, ok := strings.Cut(id, "_session_"); ok && strings.HasPrefix(id, "user_") { + return before, after + } + return "", "" +} + func prefixInvalid(prefix string, err error) error { if st, ok := status.FromError(err); ok && st.Code() == codes.InvalidArgument { return invalid("%s: %s", prefix, st.Message()) @@ -198,7 +229,7 @@ func prefixInvalid(prefix string, err error) error { // userMessages turns a user turn into its tool results, each a tool message, // then one user message with the rest. Images in tool results join that user // message: tool messages carry text only. -func userMessages(ctx context.Context, blocks []anthropicBlock) ([]*pb.ChatMessage, error) { +func userMessages(ctx context.Context, blocks []anthropicBlock, fetch bool) ([]*pb.ChatMessage, error) { var out []*pb.ChatMessage var texts []string var parts []*pb.ContentPart @@ -211,13 +242,13 @@ func userMessages(ctx context.Context, blocks []anthropicBlock) ([]*pb.ChatMessa case "text": addText(b.Text) case "image": - p, err := imagePart(ctx, b.Source) + p, err := imagePart(ctx, b.Source, fetch) if err != nil { return nil, err } parts = append(parts, p) case "document": - p, text, err := documentPart(ctx, b) + p, text, err := documentPart(ctx, b, fetch) if err != nil { return nil, err } @@ -237,7 +268,7 @@ func userMessages(ctx context.Context, blocks []anthropicBlock) ([]*pb.ChatMessa case "text": result = append(result, ib.Text) case "image": - p, err := imagePart(ctx, ib.Source) + p, err := imagePart(ctx, ib.Source, fetch) if err != nil { return nil, err } @@ -290,7 +321,7 @@ func assistantMessage(blocks []anthropicBlock) ([]*pb.ChatMessage, error) { // mediaBytes reads a base64 or https source; url sources go through safefetch // like image_url on chat completions. -func mediaBytes(ctx context.Context, src *anthropicSource, what string) (mime string, data []byte, err error) { +func mediaBytes(ctx context.Context, src *anthropicSource, what string, fetch bool) (mime string, data []byte, err error) { if src == nil { return "", nil, invalid("%s: source is required", what) } @@ -304,20 +335,23 @@ func mediaBytes(ctx context.Context, src *anthropicSource, what string) (mime st if !strings.HasPrefix(src.URL, "https://") { return "", nil, invalid("%s: url %q must be https", what, src.URL) } + if !fetch { + return "", nil, nil + } if data, mime, err = safefetch.FetchURL(ctx, src.URL, maxRequestBytes); err != nil { return "", nil, invalid("%s: failed to fetch %s: %v", what, src.URL, err) } default: return "", nil, invalid("%s: source type %q is not supported, use base64 or url", what, src.Type) } - if err := safefetch.ValidateMediaMimeType(mime); err != nil { + if err := safefetch.ValidateMediaMimeType(mime); err != nil && (fetch || src.Type != "url") { return "", nil, invalid("%s: %v", what, err) } return mime, data, nil } -func imagePart(ctx context.Context, src *anthropicSource) (*pb.ContentPart, error) { - mime, data, err := mediaBytes(ctx, src, "image") +func imagePart(ctx context.Context, src *anthropicSource, fetch bool) (*pb.ContentPart, error) { + mime, data, err := mediaBytes(ctx, src, "image", fetch) if err != nil { return nil, err } @@ -328,7 +362,7 @@ func imagePart(ctx context.Context, src *anthropicSource) (*pb.ContentPart, erro // documentPart returns a document part, or the text of a plain-text or // content-block document. -func documentPart(ctx context.Context, b anthropicBlock) (*pb.ContentPart, string, error) { +func documentPart(ctx context.Context, b anthropicBlock, fetch bool) (*pb.ContentPart, string, error) { if b.Source != nil { switch b.Source.Type { case "text": @@ -348,7 +382,7 @@ func documentPart(ctx context.Context, b anthropicBlock) (*pb.ContentPart, strin return nil, strings.Join(texts, "\n\n"), nil } } - mime, data, err := mediaBytes(ctx, b.Source, "document") + mime, data, err := mediaBytes(ctx, b.Source, "document", fetch) if err != nil { return nil, "", err } @@ -365,15 +399,11 @@ func (a *OpenAIAPI) messages(ctx context.Context, w http.ResponseWriter, r *http if req.MaxTokens < 1 { return invalid("max_tokens: must be at least 1") } - preq, err := req.toProto(ctx, getPlacementHeader(r)) + preq, err := req.toProto(ctx, getPlacementHeader(r), true) if err != nil { return err } - if req.EnableCache != nil { - preq.EnableCache = *req.EnableCache - } else if h := getCacheHeader(r); h != "" { - preq.EnableCache = h == "true" || h == "1" - } + preq.EnableCache = cacheEnabled(r, req.EnableCache) id := "msg_" + randomID() if req.Stream { @@ -476,23 +506,25 @@ func (a *OpenAIAPI) streamMessages(ctx context.Context, w http.ResponseWriter, p return err } if err == nil { - s.start(nil) - s.finish() + if err = s.start(nil); err == nil { + err = s.finish() + } + if err != nil { + log.Debug().Err(err).Msg("messages: the client left before the stream ended") + } return nil } // Headers are out: the error goes in the stream, as the Messages API sends it. - st, _ := status.FromError(err) - code, _ := httpStatus(st.Code()) - msg := st.Message() - if st.Code() == codes.Internal { - msg = "internal error" + code, _, msg, _ := apiError(w, err) + if werr := s.event("error", map[string]any{"type": "error", "error": map[string]any{"type": anthropicErrorType(code), "message": msg}}); werr != nil { + log.Debug().Err(werr).Msg("messages: the client left before the stream error") } - s.event("error", map[string]any{"type": "error", "error": map[string]any{"type": anthropicErrorType(code), "message": msg}}) return nil } // messagesStream renders GenerateText chunks as Messages SSE events. Text -// and each tool call are content blocks, opened and closed in order. +// and each tool call are content blocks, opened and closed in order; a tool +// block opens once its name is complete, at its first argument delta. type messagesStream struct { grpc.ServerStream ctx context.Context @@ -507,7 +539,8 @@ type messagesStream struct { block int // index of the last block opened open bool // whether that block is still open blockTool bool - tool int32 // vendor index of the open tool call + tool int32 // vendor index of the open or pending tool call + pending *pb.ToolCallChunk toolBlocks map[int32]bool toolCalls bool finishedAs string @@ -516,15 +549,23 @@ type messagesStream struct { func (s *messagesStream) Context() context.Context { return s.ctx } -func (s *messagesStream) event(name string, data any) { - b, _ := json.Marshal(data) - _, _ = fmt.Fprintf(s.w, "event: %s\ndata: %s\n\n", name, b) +// event writes one SSE event; its error stops the vendor stream once the +// client is gone. +func (s *messagesStream) event(name string, data any) error { + b, err := json.Marshal(data) + if err != nil { + return err + } + if _, err := fmt.Fprintf(s.w, "event: %s\ndata: %s\n\n", name, b); err != nil { + return err + } s.flusher.Flush() + return nil } -func (s *messagesStream) start(chunk *pb.GenerateTextChunk) { +func (s *messagesStream) start(chunk *pb.GenerateTextChunk) error { if s.started { - return + return nil } if chunk != nil { setRouteHeaders(s.w, chunk) @@ -533,54 +574,70 @@ func (s *messagesStream) start(chunk *pb.GenerateTextChunk) { setCacheHeader(s.w, s.enableCache, s.cacheHit) startEventStream(s.w, s.flusher) s.started = true - s.event("message_start", map[string]any{"type": "message_start", "message": map[string]any{ + if err := s.event("message_start", map[string]any{"type": "message_start", "message": map[string]any{ "id": s.id, "type": "message", "role": "assistant", "model": s.model, "content": []any{}, "stop_reason": nil, "stop_sequence": nil, "usage": anthropicUsage(nil, s.estimate), - }}) - s.event("ping", map[string]any{"type": "ping"}) + }}); err != nil { + return err + } + return s.event("ping", map[string]any{"type": "ping"}) } -func (s *messagesStream) openBlock(block map[string]any, tool bool, index int32) { - s.closeBlock() +func (s *messagesStream) openBlock(block map[string]any, tool bool) error { + if err := s.closeBlock(); err != nil { + return err + } s.block++ - s.open, s.blockTool, s.tool = true, tool, index - s.event("content_block_start", map[string]any{"type": "content_block_start", "index": s.block, "content_block": block}) + s.open, s.blockTool = true, tool + return s.event("content_block_start", map[string]any{"type": "content_block_start", "index": s.block, "content_block": block}) } -func (s *messagesStream) closeBlock() { +// closeBlock ends the open block, opening a pending tool call first so a +// call without arguments still reaches the client. +func (s *messagesStream) closeBlock() error { + if err := s.flushPending(); err != nil { + return err + } if !s.open { - return + return nil } - s.event("content_block_stop", map[string]any{"type": "content_block_stop", "index": s.block}) s.open = false + return s.event("content_block_stop", map[string]any{"type": "content_block_stop", "index": s.block}) +} + +func (s *messagesStream) flushPending() error { + p := s.pending + if p == nil { + return nil + } + s.pending = nil + return s.openBlock(map[string]any{"type": "tool_use", "id": toolUseID(p.Id), "name": p.Name, "input": map[string]any{}}, true) } -func (s *messagesStream) delta(delta map[string]any) { - s.event("content_block_delta", map[string]any{"type": "content_block_delta", "index": s.block, "delta": delta}) +func (s *messagesStream) delta(delta map[string]any) error { + return s.event("content_block_delta", map[string]any{"type": "content_block_delta", "index": s.block, "delta": delta}) } func (s *messagesStream) Send(chunk *pb.GenerateTextChunk) error { if chunk.Cached { s.cacheHit = true } - s.start(chunk) + if err := s.start(chunk); err != nil { + return err + } if chunk.Delta != "" { - if !s.open || s.blockTool { - s.openBlock(map[string]any{"type": "text", "text": ""}, false, -1) + if !s.open || s.blockTool || s.pending != nil { + if err := s.openBlock(map[string]any{"type": "text", "text": ""}, false); err != nil { + return err + } + } + if err := s.delta(map[string]any{"type": "text_delta", "text": chunk.Delta}); err != nil { + return err } - s.delta(map[string]any{"type": "text_delta", "text": chunk.Delta}) } for _, tc := range chunk.ToolCallChunks { - if !s.open || !s.blockTool || s.tool != tc.Index { - if s.toolBlocks[tc.Index] { - return status.Errorf(codes.Internal, "tool call %d continued after the next content block began", tc.Index) - } - s.toolBlocks[tc.Index] = true - s.toolCalls = true - s.openBlock(map[string]any{"type": "tool_use", "id": toolUseID(tc.Id), "name": tc.Name, "input": map[string]any{}}, true, tc.Index) - } - if tc.ArgumentsDelta != "" { - s.delta(map[string]any{"type": "input_json_delta", "partial_json": tc.ArgumentsDelta}) + if err := s.toolChunk(tc); err != nil { + return err } } if chunk.FinishReason != "" { @@ -592,13 +649,48 @@ func (s *messagesStream) Send(chunk *pb.GenerateTextChunk) error { return nil } -func (s *messagesStream) finish() { - s.closeBlock() +func (s *messagesStream) toolChunk(tc *pb.ToolCallChunk) error { + current := s.pending != nil || s.open && s.blockTool + if !current || s.tool != tc.Index { + if s.toolBlocks[tc.Index] { + return status.Errorf(codes.Internal, "tool call %d continued after the next content block began", tc.Index) + } + if err := s.closeBlock(); err != nil { + return err + } + s.toolBlocks[tc.Index] = true + s.toolCalls = true + s.tool = tc.Index + s.pending = &pb.ToolCallChunk{Index: tc.Index} + } + if p := s.pending; p != nil { + // Names arrive as deltas, like arguments. + p.Name += tc.Name + if p.Id == "" { + p.Id = tc.Id + } + } + // "null" is how a vendor spells no arguments; the block starts with {}. + if tc.ArgumentsDelta == "" || tc.ArgumentsDelta == "null" && s.pending != nil { + return nil + } + if err := s.flushPending(); err != nil { + return err + } + return s.delta(map[string]any{"type": "input_json_delta", "partial_json": tc.ArgumentsDelta}) +} + +func (s *messagesStream) finish() error { + if err := s.closeBlock(); err != nil { + return err + } usage := anthropicUsage(s.usage, s.estimate) - s.event("message_delta", map[string]any{"type": "message_delta", + if err := s.event("message_delta", map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": stopReason(s.finishedAs, s.toolCalls), "stop_sequence": nil}, - "usage": usage}) - s.event("message_stop", map[string]any{"type": "message_stop"}) + "usage": usage}); err != nil { + return err + } + return s.event("message_stop", map[string]any{"type": "message_stop"}) } func writeAnthropicError(w http.ResponseWriter, err error) { @@ -620,8 +712,6 @@ func anthropicErrorType(code int) string { return "permission_error" case http.StatusNotFound: return "not_found_error" - case http.StatusRequestEntityTooLarge: - return "request_too_large" case http.StatusTooManyRequests: return "rate_limit_error" case http.StatusServiceUnavailable: @@ -633,28 +723,17 @@ func anthropicErrorType(code int) string { } } -// countTokens estimates like the token budgets do (a rune count), adding the -// tool definitions and tool calls the estimate leaves out. +// countTokens answers with the estimate the token budgets reserve. func (a *OpenAIAPI) countTokens(ctx context.Context, w http.ResponseWriter, r *http.Request) error { var req messagesRequest if err := decodeBody(r, &req); err != nil { return err } - preq, err := req.toProto(ctx, "") + preq, err := req.toProto(ctx, "", false) if err != nil { return err } - n := grpcapi.EstimateTokens(preq) - extra := 0 - for _, t := range req.Tools { - extra += utf8.RuneCountInString(t.Name) + utf8.RuneCountInString(t.Description) + utf8.RuneCount(t.InputSchema) - } - for _, m := range preq.Messages { - for _, tc := range m.ToolCalls { - extra += utf8.RuneCountInString(tc.GetFunction().GetName()) + utf8.RuneCountInString(tc.GetFunction().GetArguments()) - } - } setObservabilityHeader(w, "Token-Count", "estimate") - writeJSON(w, http.StatusOK, map[string]any{"input_tokens": n + int64(extra/4)}) + writeJSON(w, http.StatusOK, map[string]any{"input_tokens": grpcapi.EstimateTokens(preq)}) return nil } diff --git a/internal/server/anthropic_test.go b/internal/server/anthropic_test.go index cfcbb9a..1190518 100644 --- a/internal/server/anthropic_test.go +++ b/internal/server/anthropic_test.go @@ -47,7 +47,7 @@ func TestMessagesMapsRequest(t *testing.T) { rec := messagesCall(h, "/v1/messages?beta=true", key, `{ "model":"auto:high","max_tokens":32000, "system":[{"type":"text","text":"You are Claude Code."},{"type":"text","text":"Be brief.","cache_control":{"type":"ephemeral"}}], - "metadata":{"user_id":"user_abc_session_1"}, + "metadata":{"user_id":"u-42"}, "temperature":1,"thinking":{"type":"enabled","budget_tokens":1024}, "tools":[{"name":"Read","description":"Read a file","input_schema":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}], "tool_choice":{"type":"any"}, @@ -68,7 +68,7 @@ func TestMessagesMapsRequest(t *testing.T) { if r.Effort != pb.Effort_EFFORT_HIGH || !r.AutoRoute || r.MaxOutputTokens != 32000 { t.Errorf("routing: %+v", r) } - if r.BaseSystemInstruction != "You are Claude Code.\n\nBe brief." || r.Attributes["user_id"] != "user_abc_session_1" { + if r.BaseSystemInstruction != "You are Claude Code.\n\nBe brief." || r.Attributes["user_id"] != "u-42" { t.Errorf("system %q attrs %v", r.BaseSystemInstruction, r.Attributes) } if len(r.Tools) != 1 || r.Tools[0].Function.Name != "Read" || r.Tools[0].Function.Parameters.Properties["path"].Type != "string" || @@ -267,19 +267,79 @@ func TestMessagesErrors(t *testing.T) { } } +func TestClaudeCodeUser(t *testing.T) { + for _, c := range []struct{ in, user, session string }{ + {`{"device_id":"dev1","account_uuid":"","session_id":"s1"}`, "dev1", "s1"}, + {`{"device_id":"dev1","account_uuid":"acc1","session_id":"s2"}`, "acc1", "s2"}, + {"user_abc_account_acc1_session_s3", "user_abc_account_acc1", "s3"}, + {"u-42", "", ""}, + {"{not json", "", ""}, + } { + if u, s := claudeCodeUser(c.in); u != c.user || s != c.session { + t.Errorf("%s: %q %q", c.in, u, s) + } + } + fp, h, key := newTestAPI(t) + messagesCall(h, "/v1/messages", key, `{"max_tokens":1,"metadata":{"user_id":"{\"device_id\":\"dev1\",\"account_uuid\":\"\",\"session_id\":\"s1\"}"},"messages":[{"role":"user","content":"x"}]}`) + if fp.text.Attributes["user_id"] != "dev1" || fp.text.Attributes["session_id"] != "s1" { + t.Errorf("attributes %v", fp.text.Attributes) + } +} + +func TestMessagesStreamToolDeltas(t *testing.T) { + fp, h, key := newTestAPI(t) + fp.streamChunks = []*pb.GenerateTextChunk{ + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, Id: "a", Name: "Re"}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, Name: "ad"}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 0, ArgumentsDelta: `{"path":"a.go"}`}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 1, Id: "b", Name: "Glob", ArgumentsDelta: "null"}}}, + {ToolCallChunks: []*pb.ToolCallChunk{{Index: 2, Id: "c", Name: "Ls"}}}, + {FinishReason: "tool_calls"}, + } + rec := messagesCall(h, "/v1/messages", key, `{"max_tokens":10,"stream":true,"messages":[{"role":"user","content":"x"}]}`) + b := rec.Body.String() + for _, s := range []string{ + `"content_block":{"id":"a","input":{},"name":"Read","type":"tool_use"},"index":0`, + `"content_block":{"id":"b","input":{},"name":"Glob","type":"tool_use"},"index":1`, + `"content_block":{"id":"c","input":{},"name":"Ls","type":"tool_use"},"index":2`, + `{"index":2,"type":"content_block_stop"}`, + } { + if !strings.Contains(b, s) { + t.Errorf("stream missing %s:\n%s", s, b) + } + } + if strings.Contains(b, `"partial_json":"null"`) { + t.Errorf("null arguments were streamed:\n%s", b) + } +} + +func TestMessagesStreamErrorEvent(t *testing.T) { + fp, h, key := newTestAPI(t) + fp.streamChunks = []*pb.GenerateTextChunk{{Delta: "part"}} + fp.streamErr = status.Error(codes.Unavailable, "vendor_error") + rec := messagesCall(h, "/v1/messages", key, `{"max_tokens":10,"stream":true,"messages":[{"role":"user","content":"x"}]}`) + b := rec.Body.String() + if ev := sseEvents(b); ev[len(ev)-1] != "error" || + !strings.Contains(b, `"error":{"message":"vendor_error","type":"api_error"}`) { + t.Errorf("stream error, as the non-stream 502:\n%s", b) + } +} + func TestCountTokens(t *testing.T) { _, h, key := newTestAPI(t) rec := messagesCall(h, "/v1/messages/count_tokens", key, `{"model":"auto", "system":"`+strings.Repeat("s", 400)+`", "tools":[{"name":"Read","input_schema":{"type":"object"}}], - "messages":[{"role":"user","content":"`+strings.Repeat("u", 400)+`"}]}`) + "messages":[{"role":"user","content":[{"type":"text","text":"`+strings.Repeat("u", 400)+`"}, + {"type":"image","source":{"type":"url","url":"https://127.0.0.1:1/never-fetched.png"}}]}]}`) var out struct { InputTokens int64 `json:"input_tokens"` } if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || rec.Code != 200 { t.Fatalf("%d %s", rec.Code, rec.Body) } - if out.InputTokens < 200 || out.InputTokens > 220 || rec.Header().Get("X-KTMS-Token-Count") != "estimate" { + // 800 runes, the tool, and the image's 1000. + if out.InputTokens < 1200 || out.InputTokens > 1220 || rec.Header().Get("X-KTMS-Token-Count") != "estimate" { t.Errorf("count %d headers %v", out.InputTokens, rec.Header()) } } diff --git a/internal/server/openai.go b/internal/server/openai.go index 0217779..7d94482 100644 --- a/internal/server/openai.go +++ b/internal/server/openai.go @@ -690,13 +690,7 @@ func (a *OpenAIAPI) chatCompletions(ctx context.Context, w http.ResponseWriter, if err != nil { return err } - var enableCache bool - if req.EnableCache != nil { - enableCache = *req.EnableCache - } else if h := getCacheHeader(r); h != "" { - enableCache = h == "true" || h == "1" - } - preq.EnableCache = enableCache + preq.EnableCache = cacheEnabled(r, req.EnableCache) id, created := "chatcmpl-"+randomID(), time.Now().Unix() if req.Stream { @@ -1006,6 +1000,15 @@ func getPlacementHeader(r *http.Request) string { return "" } +// cacheEnabled is the body's enable_cache, else the X-KTMS-Cache header. +func cacheEnabled(r *http.Request, body *bool) bool { + if body != nil { + return *body + } + h := getCacheHeader(r) + return h == "true" || h == "1" +} + func getCacheHeader(r *http.Request) string { if v := r.Header.Get("X-KTMS-Cache"); v != "" { return v diff --git a/internal/server/openai_test.go b/internal/server/openai_test.go index 840d1ce..8cb799d 100644 --- a/internal/server/openai_test.go +++ b/internal/server/openai_test.go @@ -35,6 +35,7 @@ type fakeProxy struct { respJudgeScore float64 respJudgeModel string streamChunks []*pb.GenerateTextChunk + streamErr error // returned after streamChunks emptyStream bool media map[string]string speechErr error @@ -79,7 +80,7 @@ func (f *fakeProxy) GenerateTextStream(r *pb.GenerateTextStreamRequest, s pb.Gat return err } } - return nil + return f.streamErr } if err := s.Send(&pb.GenerateTextChunk{ Delta: "hello",