Merge pull request 'feat: add multi-key support with automatic rate-limit rotation' (#11) from renekv/ocgo:feat/multi-key-rotation into main

Reviewed-on: #11

Makefile

@@ -1,12 +1,14 @@
 .PHONY: build run test clean install release
 
 VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev)
+EXE := $(shell go env GOEXE)
+OCGO_BIN := bin/ocgo$(EXE)
 GOBIN := $(shell go env GOBIN)
 GOPATH := $(shell go env GOPATH)
-INSTALL_DIR := $(if $(GOBIN),$(GOBIN),$(GOPATH)/bin)
+INSTALL_DIR := $(HOME)/.local/bin
 
 build:
-	go build -ldflags "-X main.version=$(VERSION)" -o bin/ocgo ./cmd/ocgo
+	go build -ldflags "-X main.version=$(VERSION)" -o $(OCGO_BIN) ./cmd/ocgo
 
 run:
 	go run ./cmd/ocgo
@@ -19,7 +21,7 @@ clean:
 
 install: build
 	mkdir -p "$(INSTALL_DIR)"
-	install -m 0755 bin/ocgo "$(INSTALL_DIR)/ocgo"
+	install -m 0755 "$(OCGO_BIN)" "$(INSTALL_DIR)/ocgo$(EXE)"
 
 release:
 	@[ -n "$(TAG)" ] || (echo "Usage: make release TAG=v0.1.0" && exit 1)

build.bat

@@ -0,0 +1,9 @@
+@echo off
+setlocal
+
+for /f "delims=" %%v in ('git describe --tags --always --dirty 2^>NUL') do set "OCGO_VERSION=%%v"
+if not defined OCGO_VERSION set "OCGO_VERSION=dev"
+
+if not exist bin mkdir bin
+go build -ldflags "-X main.version=%OCGO_VERSION%" -o bin\ocgo.exe .\cmd\ocgo
+exit /b %ERRORLEVEL%

cmd/ocgo/listener_pid_netstat.go

@@ -0,0 +1,40 @@
+package main
+
+import (
+	"errors"
+	"strconv"
+	"strings"
+)
+
+func parseWindowsNetstatPID(output string, port int) (int, error) {
+	wanted := strconv.Itoa(port)
+	for _, line := range strings.Split(output, "\n") {
+		fields := strings.Fields(line)
+		if len(fields) < 5 || !strings.EqualFold(fields[0], "tcp") {
+			continue
+		}
+		if !strings.EqualFold(fields[len(fields)-2], "listening") {
+			continue
+		}
+		if !netstatAddressUsesPort(fields[1], wanted) {
+			continue
+		}
+		pid, err := strconv.Atoi(fields[len(fields)-1])
+		if err == nil && pid > 0 {
+			return pid, nil
+		}
+	}
+	return 0, errors.New("no listener found")
+}
+
+func netstatAddressUsesPort(address, port string) bool {
+	address = strings.TrimSpace(address)
+	if address == "" {
+		return false
+	}
+	if strings.HasPrefix(address, "[") {
+		return strings.HasSuffix(address, "]:"+port)
+	}
+	idx := strings.LastIndex(address, ":")
+	return idx >= 0 && address[idx+1:] == port
+}

cmd/ocgo/listener_pid_unix.go

@@ -0,0 +1,27 @@
+//go:build !windows
+
+package main
+
+import (
+	"errors"
+	"os/exec"
+	"strconv"
+	"strings"
+)
+
+func findListenerPID(port int) (int, error) {
+	if port == 0 {
+		return 0, errors.New("missing port")
+	}
+	out, err := exec.Command("lsof", "-nP", "-tiTCP:"+strconv.Itoa(port), "-sTCP:LISTEN").Output()
+	if err != nil {
+		return 0, err
+	}
+	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
+		pid, err := strconv.Atoi(strings.TrimSpace(line))
+		if err == nil && pid > 0 {
+			return pid, nil
+		}
+	}
+	return 0, errors.New("no listener found")
+}

cmd/ocgo/listener_pid_windows.go

@@ -0,0 +1,19 @@
+//go:build windows
+
+package main
+
+import (
+	"errors"
+	"os/exec"
+)
+
+func findListenerPID(port int) (int, error) {
+	if port == 0 {
+		return 0, errors.New("missing port")
+	}
+	out, err := exec.Command("netstat", "-ano", "-p", "tcp").Output()
+	if err != nil {
+		return 0, err
+	}
+	return parseWindowsNetstatPID(string(out), port)
+}

cmd/ocgo/main_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
@@ -59,6 +60,36 @@ func TestWriteCodexProfileReplacesExistingSections(t *testing.T) {
 	}
 }
 
+func TestModelContextWindow(t *testing.T) {
+	if got := modelContextWindow("deepseek-v4-flash[1M]"); got != 1000000 {
+		t.Fatalf("deepseek-v4-flash[1M]: got %d, want 1000000", got)
+	}
+	if got := modelContextWindow("deepseek-v4-pro[1M]"); got != 1000000 {
+		t.Fatalf("deepseek-v4-pro[1M]: got %d, want 1000000", got)
+	}
+	if got := modelContextWindow("kimi-k2.6"); got != 128000 {
+		t.Fatalf("kimi-k2.6: got %d, want 128000", got)
+	}
+	if got := modelContextWindow("unknown-model"); got != 128000 {
+		t.Fatalf("unknown-model: got %d, want 128000", got)
+	}
+}
+
+func TestCleanModelName(t *testing.T) {
+	tests := []struct{ in, want string }{
+		{"deepseek-v4-flash[1M]", "deepseek-v4-flash"},
+		{"deepseek-v4-pro[1M]", "deepseek-v4-pro"},
+		{"deepseek-v4-flash", "deepseek-v4-flash"},
+		{"kimi-k2.6", "kimi-k2.6"},
+		{"kimi-k2.6[128K]", "kimi-k2.6"},
+	}
+	for _, tt := range tests {
+		if got := cleanModelName(tt.in); got != tt.want {
+			t.Errorf("cleanModelName(%q) = %q, want %q", tt.in, got, tt.want)
+		}
+	}
+}
+
 func TestWriteCodexModelCatalog(t *testing.T) {
 	path := filepath.Join(t.TempDir(), "ocgo-models.json")
 	if err := writeCodexModelCatalog(path); err != nil {
@@ -69,13 +100,34 @@ func TestWriteCodexModelCatalog(t *testing.T) {
 		t.Fatal(err)
 	}
 	content := string(b)
-	for _, want := range []string{`"models"`, `"slug": "deepseek-v4-pro"`, `"context_window": 128000`, `"truncation_policy"`} {
+	for _, want := range []string{`"models"`, `"slug": "deepseek-v4-pro[1M]"`, `"context_window": 1000000`, `"truncation_policy"`, `"supports_image_detail_original": false`, `"image"`} {
 		if !strings.Contains(content, want) {
 			t.Fatalf("missing %q in:\n%s", want, content)
 		}
 	}
 }
 
+func TestCodexModelCatalogAllowsImagesForKnownVisionModels(t *testing.T) {
+	if !modelSupportsImages("kimi-k2.6") {
+		t.Fatal("kimi-k2.6 should support image inputs")
+	}
+	if modelSupportsImages("deepseek-v4-pro") {
+		t.Fatal("deepseek-v4-pro should not support image inputs")
+	}
+	for _, tc := range []struct {
+		model string
+		want  []string
+	}{
+		{model: "kimi-k2.6", want: []string{"text", "image"}},
+		{model: "deepseek-v4-pro", want: []string{"text"}},
+	} {
+		got := modelInputModalities(tc.model)
+		if strings.Join(got, ",") != strings.Join(tc.want, ",") {
+			t.Fatalf("%s modalities = %+v, want %+v", tc.model, got, tc.want)
+		}
+	}
+}
+
 func TestCompareVersions(t *testing.T) {
 	if compareVersions("0.80.9", "0.81.0") >= 0 {
 		t.Fatal("0.80.9 should be older")
@@ -141,11 +193,248 @@ func TestAnthropicToolResultPreservesFollowingUserText(t *testing.T) {
 	if messages[0].Role != "tool" || messages[0].ToolCallID != "call_123" || messages[0].Content != "09:33:16" {
 		t.Fatalf("bad tool result conversion: %+v", messages[0])
 	}
-	if messages[1].Role != "user" || !strings.Contains(messages[1].Content, "figma.example") {
+	if messages[1].Role != "user" || !strings.Contains(contentString(messages[1].Content), "figma.example") {
 		t.Fatalf("following user text was not preserved: %+v", messages[1])
 	}
 }
 
+func TestResponsesInputPreservesImages(t *testing.T) {
+	messages := responsesInputToMessages([]byte(`[{"type":"message","role":"user","content":[{"type":"input_text","text":"describe this"},{"type":"input_image","image_url":"data:image/png;base64,abc","detail":"high"}]}]`))
+	if len(messages) != 1 {
+		t.Fatalf("got %d messages", len(messages))
+	}
+	parts, ok := messages[0].Content.([]OAIContentPart)
+	if !ok {
+		t.Fatalf("content should be multimodal parts: %+v", messages[0].Content)
+	}
+	if len(parts) != 2 || parts[0].Type != "text" || parts[0].Text != "describe this" {
+		t.Fatalf("bad text part: %+v", parts)
+	}
+	if parts[1].Type != "image_url" || parts[1].ImageURL == nil || parts[1].ImageURL.URL != "data:image/png;base64,abc" || parts[1].ImageURL.Detail != "" {
+		t.Fatalf("bad image part: %+v", parts[1])
+	}
+}
+
+func TestResponsesImageKeepsKimiModel(t *testing.T) {
+	req := ResponsesRequest{Model: "kimi-k2.6", Input: []byte(`[{"type":"message","role":"user","content":[{"type":"input_text","text":"describe this"},{"type":"input_image","image_url":"data:image/png;base64,abc"}]}]`)}
+	out := responsesToChat(req, "")
+	if out.Model != "kimi-k2.6" {
+		t.Fatalf("image request should keep Kimi model, got %q", out.Model)
+	}
+	if err := validateImageSupport(out); err != nil {
+		t.Fatalf("Kimi image request should validate: %v", err)
+	}
+}
+
+func TestResponsesImageRejectsUnsupportedModel(t *testing.T) {
+	req := ResponsesRequest{Model: "deepseek-v4-pro", Input: []byte(`[{"type":"message","role":"user","content":[{"type":"input_text","text":"describe this"},{"type":"input_image","image_url":"data:image/png;base64,abc"}]}]`)}
+	out := responsesToChat(req, "")
+	if err := validateImageSupport(out); err == nil || !strings.Contains(err.Error(), "deepseek-v4-pro") {
+		t.Fatalf("DeepSeek image request should be rejected, got %v", err)
+	}
+}
+
+func TestRawChatImageKeepsKimiAndStripsDetail(t *testing.T) {
+	body, err := prepareChatBody([]byte(`{"model":"kimi-k2.6","messages":[{"role":"user","content":[{"type":"text","text":"describe this"},{"type":"image_url","image_url":{"url":"data:image/png;base64,abc","detail":"high"}}]}]}`))
+	if err != nil {
+		t.Fatalf("Kimi image request should validate: %v", err)
+	}
+	if !strings.Contains(string(body), `"model":"kimi-k2.6"`) {
+		t.Fatalf("image chat body should keep Kimi model: %s", string(body))
+	}
+	if strings.Contains(string(body), `"detail"`) {
+		t.Fatalf("image detail should be stripped for compatibility: %s", string(body))
+	}
+}
+
+func TestRawChatImageRejectsUnsupportedModel(t *testing.T) {
+	_, err := prepareChatBody([]byte(`{"model":"deepseek-v4-pro","messages":[{"role":"user","content":[{"type":"text","text":"describe this"},{"type":"image_url","image_url":{"url":"data:image/png;base64,abc"}}]}]}`))
+	if err == nil || !strings.Contains(err.Error(), "deepseek-v4-pro") {
+		t.Fatalf("DeepSeek image request should be rejected, got %v", err)
+	}
+}
+
+func TestRawChatStreamRequestsUsage(t *testing.T) {
+	body, err := prepareChatBody([]byte(`{"model":"kimi-k2.6","stream":true,"stream_options":{"foo":"bar"},"messages":[{"role":"user","content":"hello"}]}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	var req map[string]any
+	if err := json.Unmarshal(body, &req); err != nil {
+		t.Fatal(err)
+	}
+	options, ok := req["stream_options"].(map[string]any)
+	if !ok {
+		t.Fatalf("missing stream options in %s", string(body))
+	}
+	if options["include_usage"] != true || options["foo"] != "bar" {
+		t.Fatalf("bad stream options: %+v", options)
+	}
+}
+
+func TestConvertedStreamingRequestsAskForUsage(t *testing.T) {
+	anthropic := convertRequest(AnthropicRequest{Model: "kimi-k2.6", Stream: true, Messages: []AMessage{{Role: "user", Content: []byte(`hello`)}}}, "")
+	if anthropic.StreamOptions == nil || !anthropic.StreamOptions.IncludeUsage {
+		t.Fatalf("anthropic conversion should request stream usage: %+v", anthropic.StreamOptions)
+	}
+	responses := responsesToChat(ResponsesRequest{Model: "kimi-k2.6", Stream: true, Input: []byte(`"hello"`)}, "")
+	if responses.StreamOptions == nil || !responses.StreamOptions.IncludeUsage {
+		t.Fatalf("responses conversion should request stream usage: %+v", responses.StreamOptions)
+	}
+	plain := responsesToChat(ResponsesRequest{Model: "kimi-k2.6", Input: []byte(`"hello"`)}, "")
+	if plain.StreamOptions != nil {
+		t.Fatalf("non-streaming conversion should not set stream options: %+v", plain.StreamOptions)
+	}
+}
+
+func TestAnthropicContentPreservesImages(t *testing.T) {
+	messages := contentToOpenAI(AMessage{Role: "user", Content: []byte(`[{"type":"text","text":"what is this?"},{"type":"image","source":{"type":"base64","media_type":"image/jpeg","data":"abc"}}]`)})
+	if len(messages) != 1 {
+		t.Fatalf("got %d messages", len(messages))
+	}
+	parts, ok := messages[0].Content.([]OAIContentPart)
+	if !ok {
+		t.Fatalf("content should be multimodal parts: %+v", messages[0].Content)
+	}
+	if len(parts) != 2 || parts[0].Text != "what is this?" {
+		t.Fatalf("bad text part: %+v", parts)
+	}
+	if parts[1].ImageURL == nil || parts[1].ImageURL.URL != "data:image/jpeg;base64,abc" {
+		t.Fatalf("bad image part: %+v", parts[1])
+	}
+}
+
+func TestAnthropicImageKeepsKimiModel(t *testing.T) {
+	out := convertRequest(AnthropicRequest{Model: "kimi-k2.6", Messages: []AMessage{{Role: "user", Content: []byte(`[{"type":"text","text":"what is this?"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"}}]`)}}}, "")
+	if out.Model != "kimi-k2.6" {
+		t.Fatalf("image request should keep Kimi model, got %q", out.Model)
+	}
+	if err := validateImageSupport(out); err != nil {
+		t.Fatalf("Kimi image request should validate: %v", err)
+	}
+}
+
+func TestAnthropicImageRejectsUnsupportedModel(t *testing.T) {
+	out := convertRequest(AnthropicRequest{Model: "deepseek-v4-pro", Messages: []AMessage{{Role: "user", Content: []byte(`[{"type":"text","text":"what is this?"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"}}]`)}}}, "")
+	if err := validateImageSupport(out); err == nil || !strings.Contains(err.Error(), "deepseek-v4-pro") {
+		t.Fatalf("DeepSeek image request should be rejected, got %v", err)
+	}
+}
+
+func contentString(v any) string {
+	s, _ := v.(string)
+	return s
+}
+
+func TestParseWindowsNetstatPID(t *testing.T) {
+	output := strings.Join([]string{
+		"Proto  Local Address          Foreign Address        State           PID",
+		"TCP    127.0.0.1:3456       0.0.0.0:0              LISTENING       4321",
+		"TCP    [::1]:9999           [::]:0                 LISTENING       8765",
+		"TCP    127.0.0.1:34560      0.0.0.0:0              LISTENING       1111",
+	}, "\n")
+	pid, err := parseWindowsNetstatPID(output, 3456)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if pid != 4321 {
+		t.Fatalf("pid = %d, want 4321", pid)
+	}
+}
+
+func TestParseWindowsNetstatPIDMatchesIPv6(t *testing.T) {
+	output := "TCP    [::]:3456             [::]:0                 LISTENING       2468\n"
+	pid, err := parseWindowsNetstatPID(output, 3456)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if pid != 2468 {
+		t.Fatalf("pid = %d, want 2468", pid)
+	}
+}
+
+func TestWriteAnthropicResponseIncludesUsage(t *testing.T) {
+	body := strings.NewReader(`{"choices":[{"message":{"content":"done"}}],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16,"prompt_tokens_details":{"cached_tokens":4}}}`)
+	w := httptest.NewRecorder()
+	writeAnthropicResponse(w, body, "kimi-k2.6")
+	var out map[string]any
+	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
+		t.Fatal(err)
+	}
+	usage, ok := out["usage"].(map[string]any)
+	if !ok {
+		t.Fatalf("missing usage: %+v", out)
+	}
+	if usage["input_tokens"] != float64(11) || usage["output_tokens"] != float64(5) || usage["cache_read_input_tokens"] != float64(4) {
+		t.Fatalf("bad anthropic usage: %+v", usage)
+	}
+}
+
+func TestWriteResponsesResponseIncludesUsage(t *testing.T) {
+	body := strings.NewReader(`{"choices":[{"message":{"content":"done"}}],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10,"input_tokens_details":{"cached_tokens":2}}}`)
+	w := httptest.NewRecorder()
+	writeResponsesResponse(w, body, "kimi-k2.6")
+	var out map[string]any
+	if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil {
+		t.Fatal(err)
+	}
+	usage, ok := out["usage"].(map[string]any)
+	if !ok {
+		t.Fatalf("missing usage: %+v", out)
+	}
+	if usage["input_tokens"] != float64(7) || usage["output_tokens"] != float64(3) || usage["total_tokens"] != float64(10) {
+		t.Fatalf("bad responses usage: %+v", usage)
+	}
+	details, ok := usage["input_tokens_details"].(map[string]any)
+	if !ok || details["cached_tokens"] != float64(2) {
+		t.Fatalf("bad cached details: %+v", usage["input_tokens_details"])
+	}
+}
+
+func TestParseOpenAIStreamChunkReadsUsageOnly(t *testing.T) {
+	chunk := parseOpenAIStreamChunk([]byte(`{"choices":[],"usage":{"prompt_tokens":8,"completion_tokens":4,"total_tokens":12}}`))
+	if !chunk.Usage.Present || chunk.Usage.InputTokens != 8 || chunk.Usage.OutputTokens != 4 || chunk.Usage.TotalTokens != 12 {
+		t.Fatalf("bad stream usage: %+v", chunk.Usage)
+	}
+	if chunk.Content != "" || len(chunk.ToolCalls) != 0 {
+		t.Fatalf("usage-only chunk should not include deltas: %+v", chunk)
+	}
+}
+
+func TestStreamAnthropicIncludesFinalUsage(t *testing.T) {
+	body := strings.NewReader(strings.Join([]string{
+		`data: {"choices":[{"delta":{"content":"hi"}}]}`,
+		`data: {"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}`,
+		`data: [DONE]`,
+		``,
+	}, "\n\n"))
+	w := httptest.NewRecorder()
+	streamAnthropic(w, body, "kimi-k2.6")
+	out := w.Body.String()
+	for _, want := range []string{`"input_tokens":7`, `"output_tokens":3`} {
+		if !strings.Contains(out, want) {
+			t.Fatalf("missing %q in:\n%s", want, out)
+		}
+	}
+}
+
+func TestStreamResponsesIncludesCompletedUsage(t *testing.T) {
+	body := strings.NewReader(strings.Join([]string{
+		`data: {"choices":[{"delta":{"content":"hi"}}]}`,
+		`data: {"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3,"total_tokens":10}}`,
+		`data: [DONE]`,
+		``,
+	}, "\n\n"))
+	w := httptest.NewRecorder()
+	streamResponses(w, body, "kimi-k2.6")
+	out := w.Body.String()
+	for _, want := range []string{`event: response.completed`, `"input_tokens":7`, `"output_tokens":3`, `"total_tokens":10`} {
+		if !strings.Contains(out, want) {
+			t.Fatalf("missing %q in:\n%s", want, out)
+		}
+	}
+}
+
 func TestStreamAnthropicForwardsToolCalls(t *testing.T) {
 	reasoningContentCache.Lock()
 	reasoningContentCache.byCallID = map[string]string{}