diff --git a/.claude/skills/claude-code-e2e/SKILL.md b/.claude/skills/claude-code-e2e/SKILL.md new file mode 100644 index 00000000..cf044f84 --- /dev/null +++ b/.claude/skills/claude-code-e2e/SKILL.md @@ -0,0 +1,743 @@ +--- +name: claude-code-e2e +description: > + End-to-end proof that Claude Code works turn-by-turn through Olla's Anthropic + surface against a real Ollama backend. Boots a fresh Olla in front of a + user-supplied Ollama endpoint, then exercises both legs - native Anthropic + passthrough and forced OpenAI<->Anthropic translation - with two tiers: + Tier-A direct-wire protocol assertions (curl + jq against Olla) and Tier-B a + real headless Claude Code run implementing a fixed feature in a fresh clone of + thushan/smash. Use to validate the Anthropic translation/passthrough layer + under genuine Claude Code traffic. Needs a working Ollama and the claude CLI; + not CI-safe (real LLM, real network). Runs on Sonnet. +argument-hint: "[--ollama-endpoint URL] [--ollama-model NAME] [--repo URL] [--repo-sha SHA] [--skip-tier-b] [--profile]" +model: sonnet +--- + +# /claude-code-e2e - Claude Code through Olla, end to end + +Proves the thing unit tests cannot: that a real Claude Code instance completes +multi-turn tool loops through Olla's Anthropic endpoints, against a real model, +on both the passthrough and translation paths. + +Two models are in play - keep them straight: +- **System under test**: a fresh `claude` CLI -> Olla -> Ollama -> the local + model. This is what we are proving works. +- **This orchestrator**: Sonnet (`model: sonnet`), drives the phases, runs the + curl/jq assertions, writes the report. Never delegate the whole run to one + subagent - orchestration stays here. + +## Two legs, two tiers + +| | Passthrough leg | Translation leg | +|---|---|---| +| Olla config | `passthrough_enabled: true` | `passthrough_enabled: false` | +| Path exercised | native Anthropic forwarded to Ollama `/v1/messages` | Anthropic -> OpenAI -> Ollama `/v1/chat/completions` | + +| Tier | What | Gates? | +|---|---|---| +| A - protocol | direct curl + jq against Olla's `/olla/anthropic/*` | **yes** | +| B - agent | headless Claude Code implements a fixed feature in a smash clone | no (model-dependent; reported) | + +Run order per invocation: fitness preflight -> build Olla -> clone smash -> +**leg A (passthrough)**: Tier-A then Tier-B -> reset clone, restart Olla -> +**leg B (translation)**: Tier-A then Tier-B -> report. + +## Arguments and defaults + +- `--ollama-endpoint URL` - default `http://localhost:11434`. If that is not + reachable and no flag was given, **ask the user** (AskUserQuestion) for an + endpoint rather than failing silently. +- `--ollama-model NAME` - only used if no model is already loaded. The skill + prefers the first model reported loaded by `/api/ps`. +- `--repo URL` - default `https://github.com/thushan/smash`. +- `--repo-sha SHA` - default `057ad782490b2162d77923316a623101666eea1a` (main). +- `--skip-tier-b` - run only the Tier-A protocol gate (fast). +- `--profile` - boot Olla with pprof and run the performance phase (CPU/allocs + profile of Olla's hot path under load + per-leg goroutine/heap leak check). + Off by default; numbers are informational and never gate (leak growth -> WARN). + +The Tier-B task is fixed and defined in `tasks/totalsize.md` (relative to this +skill file). Read it before leg A. + +## Output discipline (per the skills spec) + +No `set -e` - keep gathering evidence past failures. `curl -sf` for quiet +success, capture full body on failure. **Every JSON assertion goes through `jq` +- never `grep` a JSON body.** Status codes are categorised, not "non-200 = +fail": 5xx / connection error (000) -> FAIL; a 404 on `count_tokens` -> FAIL +(it is a regression, see Tier-A); model declining to call a tool -> WARN (model +capability, not an Olla bug). Echo `PASS:` / `FAIL:` / `WARN:` on every +assertion so the run is scannable. + +## Phase 0 - Setup, OS detection, cleanup trap + +Anchor at repo root and stay anchored in every phase. Forward slashes +throughout, no absolute paths. + +```bash +ROOT="$(git rev-parse --show-toplevel)"; cd "$ROOT" + +# --- OS portability: the only OS-specific bits are the Go binary suffix and +# the kill command. `go env GOEXE` yields the suffix on every platform; kill_pid +# tries POSIX kill then Windows taskkill - so no $OSTYPE branching is needed and +# nothing OS-specific is hard-coded below. --- +EXE=$(go env GOEXE) # ".exe" on Windows, "" elsewhere +kill_pid() { kill "$1" 2>/dev/null || taskkill //F //PID "$1" 2>/dev/null || true; } +TMPL="$ROOT/.claude/skills/claude-code-e2e/config.regression.yaml.tmpl" + +# --- args (defaults, then parse the user's flags over them) --- +OLLAMA="${OLLAMA_ENDPOINT:-http://localhost:11434}" # --ollama-endpoint +OLLAMA_MODEL_FLAG="" # --ollama-model +REPO="https://github.com/thushan/smash" # --repo +REPO_SHA="057ad782490b2162d77923316a623101666eea1a" # --repo-sha +SKIP_TIER_B=0 +PROFILE=0 # --profile: pprof Olla + perf phase +MAXT=180 # curl --max-time on model-path calls (s) +PROMPT_FILE="$ROOT/.claude/skills/claude-code-e2e/tasks/totalsize.prompt.txt" +while [ "$#" -gt 0 ]; do + case "$1" in + --ollama-endpoint) OLLAMA="$2"; shift 2 ;; + --ollama-model) OLLAMA_MODEL_FLAG="$2"; shift 2 ;; + --repo) REPO="$2"; shift 2 ;; + --repo-sha) REPO_SHA="$2"; shift 2 ;; + --skip-tier-b) SKIP_TIER_B=1; shift ;; + --profile) PROFILE=1; shift ;; + *) echo "WARN: ignoring unknown arg: $1"; shift ;; + esac +done + +OLLA_PORT=41151 +BASE="http://127.0.0.1:${OLLA_PORT}/olla/anthropic" +TOKEN="olla-regression-token" # dummy bearer; Olla->Ollama needs no auth +PPROF="http://127.0.0.1:19841/debug/pprof" # Olla's pprof server (hardcoded addr, on with -profile) + +RUN_TS=$(date -u +%Y%m%d-%H%M%S) +GIT_SHA=$(git rev-parse --short HEAD) +RESULTS=test/results +# NB: not named TMP - that collides with the Windows %TMP% env var and would be +# inherited by `go test`/`go build` as a bogus relative temp dir. +REGRESSION_TMP=test/regression/tmp +RUNDIR="${REGRESSION_TMP}/run-${RUN_TS}" +CLONE="${RUNDIR}/smash" +LOGDIR="${RESULTS}/logs/claude-code-e2e-${RUN_TS}" +REPORT="${RESULTS}/claude-code-e2e-${RUN_TS}.md" +mkdir -p "$RUNDIR" "$LOGDIR" "$RESULTS" + +PASS_COUNT=0 FAIL_COUNT=0 WARN_COUNT=0 +WARNS="" +assert_pass() { echo "PASS: $1"; PASS_COUNT=$((PASS_COUNT + 1)); } +assert_fail() { echo "FAIL: $1"; FAIL_COUNT=$((FAIL_COUNT + 1)); } +assert_warn() { # build WARNS as a string (no Bash-4 arrays); avoid a leading blank line + echo "WARN: $1"; WARN_COUNT=$((WARN_COUNT + 1)) + [ -n "$WARNS" ] && WARNS="${WARNS}\n - $1" || WARNS=" - $1" +} + +OLLA_PID="" +cleanup() { + echo "--- claude-code-e2e cleanup ---" + if [ -n "$OLLA_PID" ]; then + kill_pid "$OLLA_PID" && echo "INFO: killed Olla PID ${OLLA_PID}" || echo "WARN: could not kill Olla PID ${OLLA_PID}" + fi + wait 2>/dev/null || true + # Throwaway clone + rendered configs go; report + logs are preserved. + rm -rf "$RUNDIR" 2>/dev/null || true + echo "PASS: cleanup complete -- port ${OLLA_PORT} released, clone removed (preserved: ${REPORT}, ${LOGDIR})" + OVERALL=$([ "$FAIL_COUNT" -gt 0 ] && echo FAIL || echo PASS) + printf '%s %s | /claude-code-e2e | %s | %dP/%dF/%dW | %s | report:%s\n' \ + "${RUN_TS:0:8}" "${RUN_TS:9:6}" "$OVERALL" "$PASS_COUNT" "$FAIL_COUNT" "$WARN_COUNT" "$GIT_SHA" "$REPORT" \ + >> "${RESULTS}/last-runs.md" +} +trap 'cleanup' EXIT INT TERM +``` + +If port `41151` is already bound, kill the stale Olla from a previous run or +abort with a clear message before continuing. + +## Phase 1 - Required tools + fitness preflight (hard gate) + +Required tools - `command -v` each, FAIL with an install hint if missing: +`go`, `git`, `jq`, `curl`, and **`claude`** (the CLI under test - if absent, +abort: "Claude Code CLI not on PATH; install it before running this skill"). + +Then the **fitness gate** - this is your "stop the run if it doesn't meet the +rigour" check. Abort (exit non-zero so the trap cleans up) on any hard fail: + +```bash +# 1) endpoint reachable; if not and no --ollama-endpoint was given, ASK the user. +if ! curl -sf --connect-timeout 5 "${OLLAMA}/api/version" >/dev/null; then + echo "FAIL: Ollama not reachable at ${OLLAMA}" # -> AskUserQuestion for an endpoint, retry once + exit 1 +fi + +# 2) choose the model: first loaded model wins, else the --ollama-model flag. +PS=$(curl -sf "${OLLAMA}/api/ps") +MODEL=$(echo "$PS" | jq -r '.models[0].name // empty') +if [ -z "$MODEL" ]; then + [ -n "$OLLAMA_MODEL_FLAG" ] || { echo "FAIL: no model loaded and --ollama-model not given"; exit 1; } + MODEL="$OLLAMA_MODEL_FLAG" + # warm-load it so /api/ps reports its real context_length + curl -sf "${OLLAMA}/api/generate" -d "{\"model\":\"${MODEL}\",\"prompt\":\"hi\",\"stream\":false}" >/dev/null + PS=$(curl -sf "${OLLAMA}/api/ps") +fi + +# 3) tool capability is mandatory - Tier-B and Tier-A tool checks need it. +# Prefer /api/tags; fall back to /api/show (capabilities surface there too on +# newer Ollama) before hard-failing a model as non-tool-capable. +CAPS=$(curl -sf "${OLLAMA}/api/tags" | jq -r --arg m "$MODEL" '.models[] | select(.name==$m) | .capabilities[]?') +if ! echo "$CAPS" | grep -qx tools; then + CAPS=$(curl -sf "${OLLAMA}/api/show" -d "{\"name\":\"${MODEL}\"}" | jq -r '.capabilities[]?' 2>/dev/null) +fi +if echo "$CAPS" | grep -qx tools; then + assert_pass "model ${MODEL} advertises tools capability" +else + assert_fail "model ${MODEL} has no tools capability (caps: $(echo "$CAPS" | tr '\n' ',')) - cannot exercise tool loops" + exit 1 +fi +# Thinking models let A7 verify reasoning survives the translation path. +MODEL_THINKS=0; echo "$CAPS" | grep -qx thinking && MODEL_THINKS=1 + +# 4) context window - the killer constraint. /api/ps reports context_length for +# the loaded model on Ollama 0.30+; fall back to /api/show for older servers. +CTX=$(echo "$PS" | jq -r --arg m "$MODEL" '.models[] | select(.name==$m) | .context_length // empty') +if [ -z "$CTX" ]; then + CTX=$(curl -sf "${OLLAMA}/api/show" -d "{\"name\":\"${MODEL}\"}" \ + | jq -r '[.model_info[] ] as $_ | (.model_info | to_entries[] | select(.key|endswith(".context_length")) | .value) // empty' 2>/dev/null | head -1) +fi +CTX=${CTX:-0} +if [ "$CTX" -lt 16384 ]; then + assert_fail "context_length ${CTX} < 16384 - too small for Claude Code's system prompt + tools; raise OLLAMA_CONTEXT_LENGTH and reload" + exit 1 +elif [ "$CTX" -lt 32768 ]; then + assert_warn "context_length ${CTX} (<32768) - Tier-B may truncate on larger files" +else + assert_pass "context_length ${CTX} sufficient" +fi +echo "INFO: model=${MODEL} ctx=${CTX} endpoint=${OLLAMA}" +``` + +## Phase 2 - Build a fresh Olla + +```bash +# Log to a file and check the exit directly - piping to tee would mask a build +# failure behind tee's exit status and leave a stale/absent binary. +if go build -o "build/regression/olla${EXE}" . > "${LOGDIR}/build.log" 2>&1; then + assert_pass "fresh Olla built" +else + assert_fail "Olla build failed"; tail -20 "${LOGDIR}/build.log"; exit 1 +fi +``` + +## Phase 3 - Clone smash at the pinned SHA + +```bash +git clone --quiet "$REPO" "$CLONE" || { assert_fail "clone $REPO failed"; exit 1; } +git -C "$CLONE" checkout --quiet "$REPO_SHA" || { assert_fail "checkout $REPO_SHA failed"; exit 1; } +( cd "$CLONE" && go mod download ) > "${LOGDIR}/go-mod-download.log" 2>&1 \ + || { assert_fail "go mod download failed"; tail -20 "${LOGDIR}/go-mod-download.log"; exit 1; } +# baseline must be green before we let an agent touch it +( cd "$CLONE" && go test ./pkg/analysis/... ) >"${LOGDIR}/baseline-test.log" 2>&1 \ + && assert_pass "smash baseline pkg/analysis tests green at ${REPO_SHA}" \ + || { assert_fail "smash baseline tests not green - bad SHA or toolchain"; exit 1; } +``` + +## Phase 4 - Render config + run each leg + +Define two reusable helpers, then call them once per leg. + +### Render + boot Olla for a leg + +```bash +render_config() { # $1 = passthrough (true|false) -> echoes rendered path + local pt="$1" out="${RUNDIR}/config.${1}.yaml" + sed -e "s|__OLLA_PORT__|${OLLA_PORT}|g" \ + -e "s|__OLLAMA_URL__|${OLLAMA}|g" \ + -e "s|__PASSTHROUGH__|${pt}|g" \ + "$TMPL" > "$out" + echo "$out" +} + +boot_olla() { # $1 = rendered config path + # -profile turns on Olla's pprof server (localhost:19841) when --profile is set. + local prof=""; [ "$PROFILE" = 1 ] && prof="-profile" + build/regression/olla${EXE} $prof --config "$1" > "${LOGDIR}/olla-${LEG}.log" 2>&1 & + OLLA_PID=$! + # Readiness: Olla healthy, then model discovery has populated the Anthropic + # models list. Assert the list is non-empty rather than matching an exact id - + # the unifier may surface the model under an alias, so a name match is brittle. + local ok=0 i=0 + while [ "$i" -lt 60 ]; do + if curl -sf "http://127.0.0.1:${OLLA_PORT}/internal/health" >/dev/null \ + && curl -sf "${BASE}/v1/models" | jq -e '(.data | length) > 0' >/dev/null 2>&1; then + ok=1; break + fi + i=$((i + 1)); sleep 1 + done + [ "$ok" = 1 ] || { assert_fail "[${LEG}] Olla not ready in 60s (see ${LOGDIR}/olla-${LEG}.log)"; return 1; } + if [ "$PROFILE" = 1 ]; then + local j=0; while [ "$j" -lt 15 ]; do curl -sf --max-time 3 "${PPROF}/" >/dev/null 2>&1 && break; j=$((j + 1)); sleep 1; done + fi + assert_pass "[${LEG}] Olla up and models discovered" +} + +stop_olla() { [ -n "$OLLA_PID" ] && kill_pid "$OLLA_PID"; OLLA_PID=""; wait 2>/dev/null || true; } +``` + +**Execution note (read carefully):** run the whole per-leg driver (the +`for LEG ...` loop below, with the helper functions sourced ahead of it) as a +single **foreground** Bash tool call. Do *not* use the Bash tool's +`run_in_background` for Olla - backgrounding happens *inside* the script via the +`&` in `boot_olla`, so `OLLA_PID=$!` captures the pid in the same shell where +`stop_olla` and the cleanup trap can reach it. `run_in_background` would spawn a +detached shell and the pid would be lost, leaving Olla orphaned between legs. + +### Tier-A - direct-wire protocol assertions + +`tier_a` runs against `$BASE` for the current `$LEG`; `$PT` is the leg's +passthrough flag. The path each leg claims to exercise is **proved**, not +assumed, by two checks: **A2b** asserts the `X-Olla-Mode` response header +(`passthrough` present on the passthrough leg, absent on translation), and +**A6** asserts the matching `/internal/stats/translators` counter advanced over +the leg. Together they make "I flipped the config" verifiable rather than +implicit. The `count_tokens` check (A5) is the same on both legs - it is a +separate translator handler, not the message route. + +```bash +# post: write the JSON body to a file, echo ONLY the HTTP code. Avoids any +# body/code string-splitting (Go encodes a trailing newline on every response). +# usage: code=$(post URL BODY OUTFILE); then jq over OUTFILE. +post() { # $1=url $2=body $3=outfile -> echoes http_code; response headers -> $3.hdr + curl -s -o "$3" -D "${3}.hdr" -w '%{http_code}' --connect-timeout 5 --max-time "$MAXT" \ + -X POST "$1" -H "content-type: application/json" -H "x-api-key: ${TOKEN}" \ + -H "anthropic-version: 2023-06-01" -d "$2" +} + +tier_a() { + local code f + local STATS="http://127.0.0.1:${OLLA_PORT}/internal/stats/translators" + + # Snapshot translator routing counters before any message traffic, so A6 can + # prove which path actually executed (not just "I flipped the config"). + local sb=$(curl -sf --max-time 10 "$STATS") + local PB=$(echo "$sb" | jq -r '.summary.total_passthrough // 0') + local TB=$(echo "$sb" | jq -r '.summary.total_translations // 0') + + # A1 - models listing (non-empty data array) + f="${LOGDIR}/${LEG}-models.json" + code=$(curl -s -o "$f" -w '%{http_code}' --connect-timeout 5 --max-time 30 "${BASE}/v1/models" -H "x-api-key: ${TOKEN}") + if [ "$code" = 200 ] && jq -e '(.data | length) > 0' "$f" >/dev/null 2>&1; then + assert_pass "[${LEG}] A1 /v1/models lists $(jq -r '.data | length' "$f") model(s)" + else + assert_fail "[${LEG}] A1 /v1/models (code ${code})" + fi + + # A2 - non-streaming message: well-formed envelope + f="${LOGDIR}/${LEG}-nonstream.json" + # max_tokens=256: thinking models exhaust a 64-token budget inside reasoning before + # emitting any text block; 256 clears the reasoning phase on a trivial prompt. + code=$(post "${BASE}/v1/messages" \ + "{\"model\":\"${MODEL}\",\"max_tokens\":256,\"messages\":[{\"role\":\"user\",\"content\":\"Reply with exactly one word: pong\"}]}" "$f") + # Use select() rather than [0] so thinking-model responses (where content[0].type=="thinking") + # don't false-fail; we require at least one text block, not that the first block is text. + if [ "$code" = 200 ] \ + && [ "$(jq -r '.type' "$f")" = message ] \ + && [ "$(jq -r '[.content[]? | select(.type=="text")] | length' "$f")" -ge 1 ] \ + && [ "$(jq -r '.stop_reason // empty' "$f")" != "" ] \ + && [ "$(jq -r '.usage.output_tokens // 0' "$f")" -gt 0 ]; then + assert_pass "[${LEG}] A2 non-stream message well-formed (stop_reason=$(jq -r '.stop_reason' "$f"))" + elif [ "${code:0:1}" = 5 ] || [ "$code" = 000 ]; then + assert_fail "[${LEG}] A2 non-stream upstream/proxy error (code ${code})" + else + assert_fail "[${LEG}] A2 non-stream malformed (code ${code}); body in $f" + fi + + # A2b - X-Olla-Mode header proves the path. Olla sets it to 'passthrough' only + # on the passthrough route; it is absent on translation. (Header grep, not JSON.) + local mode=$(grep -i '^x-olla-mode:' "${f}.hdr" 2>/dev/null | tr -d '\r' | awk '{print $2}') + if [ "$PT" = true ]; then + [ "$mode" = passthrough ] && assert_pass "[${LEG}] A2b X-Olla-Mode: passthrough present" \ + || assert_fail "[${LEG}] A2b expected X-Olla-Mode: passthrough, got '${mode:-}'" + else + [ -z "$mode" ] && assert_pass "[${LEG}] A2b X-Olla-Mode absent (translation path)" \ + || assert_fail "[${LEG}] A2b expected no X-Olla-Mode on translation, got '${mode}'" + fi + + # A3 - streaming SSE event order + no mid-stream error event + # max_tokens=512: thinking models (e.g. gemma4) spend many tokens on reasoning before + # emitting visible content; 64 exhausts the budget during thinking and produces empty output. + curl -sN --connect-timeout 5 --max-time "$MAXT" -X POST "${BASE}/v1/messages" \ + -H "content-type: application/json" -H "x-api-key: ${TOKEN}" -H "anthropic-version: 2023-06-01" \ + -d "{\"model\":\"${MODEL}\",\"max_tokens\":512,\"stream\":true,\"messages\":[{\"role\":\"user\",\"content\":\"Count: one two three\"}]}" \ + > "${LOGDIR}/${LEG}-stream.sse" 2>/dev/null + local f="${LOGDIR}/${LEG}-stream.sse" + local n_start=$(grep -c '^event: message_start' "$f") + local n_delta=$(grep -c '^event: content_block_delta' "$f") + local n_stop=$(grep -c '^event: message_stop' "$f") + local n_err=$(grep -c '^event: error' "$f") + if [ "$n_start" -ge 1 ] && [ "$n_delta" -ge 1 ] && [ "$n_stop" -ge 1 ] && [ "$n_err" -eq 0 ]; then + assert_pass "[${LEG}] A3 SSE sequence intact (start/${n_start} delta/${n_delta} stop/${n_stop})" + elif [ "$n_err" -gt 0 ]; then + assert_fail "[${LEG}] A3 mid-stream 'event: error' x${n_err} (panic-mid-stream regression?); see ${f}" + else + assert_fail "[${LEG}] A3 SSE sequence broken (start/${n_start} delta/${n_delta} stop/${n_stop})" + fi + + # A4 - tool_use round-trip. Model declining is a WARN, not a FAIL. + f="${LOGDIR}/${LEG}-tool.json" + code=$(post "${BASE}/v1/messages" \ + "{\"model\":\"${MODEL}\",\"max_tokens\":256,\"tools\":[{\"name\":\"get_weather\",\"description\":\"Get weather for a city\",\"input_schema\":{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\"}},\"required\":[\"city\"]}}],\"messages\":[{\"role\":\"user\",\"content\":\"Use the get_weather tool for Sydney.\"}]}" "$f") + if [ "$code" != 200 ]; then + assert_fail "[${LEG}] A4 tool request failed (code ${code}); see $f" + elif [ "$(jq -r '[.content[]? | select(.type=="tool_use")] | length' "$f")" -ge 1 ]; then + if [ "$(jq -r '[.content[]? | select(.type=="tool_use")][0].id // empty' "$f")" != "" ] \ + && [ "$(jq -r '[.content[]? | select(.type=="tool_use")][0].name' "$f")" = get_weather ] \ + && [ "$(jq -r '[.content[]? | select(.type=="tool_use")][0].input | type=="object"' "$f")" = true ]; then + assert_pass "[${LEG}] A4 tool_use block well-formed (stop_reason=$(jq -r '.stop_reason' "$f"))" + else + assert_fail "[${LEG}] A4 tool_use block malformed; see $f" + fi + else + assert_warn "[${LEG}] A4 model did not emit a tool_use block (model capability, not Olla)" + fi + + # A5 - count_tokens. This is a separately-registered translator handler, not + # the /v1/messages passthrough route, so it proves Claude Code's token-count + # compatibility (Claude Code calls it for real) rather than message routing. + # It must answer, not 404 - Ollama itself declares token_counting_404, so a + # 404 here means the dedicated handler isn't intercepting. + f="${LOGDIR}/${LEG}-count.json" + code=$(post "${BASE}/v1/messages/count_tokens" \ + "{\"model\":\"${MODEL}\",\"messages\":[{\"role\":\"user\",\"content\":\"how many tokens is this\"}]}" "$f") + if [ "$code" = 404 ]; then + assert_fail "[${LEG}] A5 count_tokens 404 - limitation filter not intercepting (Ollama 404 leaked through)" + elif [ "$code" = 200 ] && [ "$(jq -r '.input_tokens // 0' "$f")" -gt 0 ]; then + assert_pass "[${LEG}] A5 count_tokens -> input_tokens=$(jq -r '.input_tokens' "$f")" + else + assert_fail "[${LEG}] A5 count_tokens unexpected (code ${code}); see $f" + fi + + # A6 - routing proof via translator stats delta. The /v1/messages calls above + # must have incremented the leg's own counter; if the wrong counter moved, the + # config flip did not actually change the path. + local sa=$(curl -sf --max-time 10 "$STATS") + local PA=$(echo "$sa" | jq -r '.summary.total_passthrough // 0') + local TA=$(echo "$sa" | jq -r '.summary.total_translations // 0') + if [ "$PT" = true ]; then + [ "$PA" -gt "$PB" ] && assert_pass "[${LEG}] A6 passthrough counter advanced (+$((PA - PB)))" \ + || assert_fail "[${LEG}] A6 no passthrough delta (before=${PB} after=${PA}); leg may have translated" + else + [ "$TA" -gt "$TB" ] && assert_pass "[${LEG}] A6 translation counter advanced (+$((TA - TB)))" \ + || assert_fail "[${LEG}] A6 no translation delta (before=${TB} after=${TA}); leg may have passed through" + fi + + # A7 - reasoning survives translation. Only meaningful on the translation leg + # with a thinking-capable model: confirms the OpenAI reasoning/reasoning_content + # field is mapped to Anthropic thinking blocks rather than dropped. WARN (not + # FAIL) if absent - the model may not emit reasoning for a trivial prompt - so + # model non-determinism never fails the gate; the unit tests are the hard guard. + if [ "$PT" = false ] && [ "${MODEL_THINKS:-0}" = 1 ]; then + if grep -q 'thinking_delta' "${LOGDIR}/${LEG}-stream.sse" 2>/dev/null; then + assert_pass "[${LEG}] A7 reasoning preserved as thinking blocks through translation" + else + assert_warn "[${LEG}] A7 no thinking blocks in translated stream (model may not have reasoned; see ${LEG}-stream.sse)" + fi + fi +} +``` + +### Tier-B - real Claude Code run (non-gating) + +Read `tasks/totalsize.md` for the prompt and assertion. Reset the clone to the +pinned baseline first so each leg starts identical. + +```bash +tier_b() { + [ "$SKIP_TIER_B" = 1 ] && { assert_warn "[${LEG}] Tier-B skipped (--skip-tier-b)"; return 0; } + git -C "$CLONE" reset --hard --quiet "$REPO_SHA" && git -C "$CLONE" clean -xfd --quiet + + # Single source of truth - the prompt lives in tasks/totalsize.prompt.txt so it + # cannot drift from the documented task. (totalsize.md describes the assertions.) + local PROMPT=$(cat "$PROMPT_FILE") + + # Fresh, non-interactive Claude Code, all traffic pinned to Olla. + ( cd "$CLONE" && \ + ANTHROPIC_BASE_URL="$BASE" \ + ANTHROPIC_AUTH_TOKEN="$TOKEN" \ + ANTHROPIC_MODEL="$MODEL" \ + ANTHROPIC_SMALL_FAST_MODEL="$MODEL" \ + DISABLE_AUTOUPDATER=1 DISABLE_TELEMETRY=1 DISABLE_ERROR_REPORTING=1 \ + CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 \ + claude -p "$PROMPT" --model "$MODEL" --output-format stream-json --verbose \ + --dangerously-skip-permissions --max-turns 40 \ + ) > "${LOGDIR}/claude-${LEG}.jsonl" 2>"${LOGDIR}/claude-${LEG}.err" + + # Capture what the agent actually did BEFORE the clone is torn down, into the + # preserved log dir, so a failed/odd run can be inspected after the fact: + # - a human-readable transcript (assistant text, tool calls, tool results, result) + # - the full diff vs baseline, including any new files (non-invasive: no staging, + # so the B3/B4 porcelain checks below still see the working-tree changes). + jq -r ' + if .type=="assistant" then (.message.content[]? | + if .type=="text" then "ASSISTANT: " + (.text // "") + elif .type=="tool_use" then "TOOL_USE[" + (.name // "?") + "]: " + ((.input // {}) | tostring) + else empty end) + elif .type=="user" then (.message.content[]? | + if .type=="tool_result" then "TOOL_RESULT" + (if .is_error == true then "(ERR)" else "" end) + ": " + + ((.content) | if type=="array" then (map(.text // "") | join(" ")) else ((. // "") | tostring) end) + else empty end) + elif .type=="result" then "--- RESULT is_error=" + (.is_error|tostring) + " ---\n" + (.result // "") + else empty end + ' "${LOGDIR}/claude-${LEG}.jsonl" > "${LOGDIR}/tierb-transcript-${LEG}.txt" 2>/dev/null + { git -C "$CLONE" --no-pager diff "$REPO_SHA"; + for nf in $(git -C "$CLONE" ls-files --others --exclude-standard); do + echo "=== NEW FILE: ${nf} ==="; cat "${CLONE}/${nf}"; done + } > "${LOGDIR}/tierb-diff-${LEG}.patch" 2>/dev/null + + # B1 - the child reported success (jq-select the result event, not grep) + local result=$(jq -c 'select(.type=="result")' "${LOGDIR}/claude-${LEG}.jsonl" 2>/dev/null | tail -1) + # .is_error is a JSON boolean; jq's // operator treats false as falsy so + # `false // true` returns true - use an explicit if/else to avoid the trap. + local is_err=$(echo "$result" | jq -r 'if .is_error == false then "false" else "true" end') + [ "$is_err" = false ] && assert_pass "[${LEG}] B1 Claude Code reported success" \ + || assert_warn "[${LEG}] B1 Claude Code did not report success (is_error=${is_err}) - model capability; see claude-${LEG}.jsonl" + + # B2 - tests green + if ( cd "$CLONE" && go test ./pkg/analysis/... ) >"${LOGDIR}/tierb-test-${LEG}.log" 2>&1; then + assert_pass "[${LEG}] B2 go test ./pkg/analysis/... green after edit" + else + assert_warn "[${LEG}] B2 tests not green after edit (model capability); see tierb-test-${LEG}.log" + fi + + # B3 - a real change landed on summary.go (porcelain catches both modified and + # newly-added paths, unlike a commit-to-worktree diff). + if git -C "$CLONE" status --porcelain | grep -q 'pkg/analysis/summary.go'; then + assert_pass "[${LEG}] B3 summary.go modified by the agent" + else + assert_warn "[${LEG}] B3 no change to summary.go (agent produced nothing usable)" + fi + + # B4 - the requested TEST was actually added (not just the method). Guards the + # B2/B3 hole where a model adds the method, skips the test, and rides the + # pre-existing suite to green. (grep on a Go source file, not JSON.) + if git -C "$CLONE" status --porcelain | grep -q 'pkg/analysis/summary_test.go' \ + && grep -q 'TotalSize' "${CLONE}/pkg/analysis/summary_test.go" 2>/dev/null; then + assert_pass "[${LEG}] B4 a TotalSize test was added" + else + assert_warn "[${LEG}] B4 no TotalSize test added (method may be untested)" + fi + + echo "INFO: [${LEG}] Tier-B artifacts in ${LOGDIR}: tierb-transcript-${LEG}.txt, tierb-diff-${LEG}.patch ($(grep -c '^' "${LOGDIR}/tierb-diff-${LEG}.patch" 2>/dev/null) diff lines), claude-${LEG}.jsonl" +} +``` + +### Performance profiling (optional, `--profile`) + +Gauges Olla's own hot path under load, *after* functional validation has passed - +numbers never gate the run (a goroutine leak is the only WARN). Two ideas: + +1. **Two burst modes per leg.** `perf_burst micro` floods `count_tokens` + + `/v1/models` - served entirely by Olla (estimator, translation envelope, + unifier/registry) with **no backend inference wait** - so it isolates Olla's + own CPU. `perf_burst stream` drives real streaming `/v1/messages` at low + concurrency to profile the **realistic** translation/proxy + SSE path + (backend-bound, so fewer requests, but it captures the true streaming + allocation profile, not the micro-request one). Compare the two: micro + over-weights per-request middleware overhead; stream shows what real Claude + Code traffic actually allocates. +2. **Bracket each leg with goroutine/heap snapshots** to catch a streaming/proxy + goroutine leak across a realistic session (incl. the Tier-B multi-turn run). + +Raw `.pb.gz` profiles are saved in `$LOGDIR` for interactive `go tool pprof`; the +headline (top cumulative CPU, top allocators) is auto-extracted into text. + +```bash +pprof_goroutines() { curl -sf --max-time 5 "${PPROF}/goroutine?debug=1" | head -1 | grep -oE '[0-9]+' | head -1; } +pprof_heap_field() { # $1 = HeapAlloc | NumGC + curl -sf --max-time 5 "${PPROF}/heap?debug=1" | grep -E "^#[[:space:]]+$1[[:space:]]*=" | head -1 | grep -oE '[0-9]+' | head -1 +} + +perf_request() { # one request appropriate to the burst mode (subshells inherit this fn) + case "$1" in + micro) # pure-Olla endpoints, no backend inference wait - isolates Olla's CPU + curl -s -o /dev/null --max-time 10 -X POST "${BASE}/v1/messages/count_tokens" \ + -H "content-type: application/json" -H "x-api-key: ${TOKEN}" -H "anthropic-version: 2023-06-01" \ + -d "{\"model\":\"${MODEL}\",\"messages\":[{\"role\":\"user\",\"content\":\"profile the hot path with a body large enough to exercise tokenisation and the translation envelope build\"}]}" + curl -s -o /dev/null --max-time 10 "${BASE}/v1/models" -H "x-api-key: ${TOKEN}" ;; + stream) # real streaming messages - exercises the true translation/proxy + SSE path + curl -sN -o /dev/null --max-time 30 -X POST "${BASE}/v1/messages" \ + -H "content-type: application/json" -H "x-api-key: ${TOKEN}" -H "anthropic-version: 2023-06-01" \ + -d "{\"model\":\"${MODEL}\",\"max_tokens\":128,\"stream\":true,\"messages\":[{\"role\":\"user\",\"content\":\"Briefly list three colours.\"}]}" ;; + esac +} + +perf_burst() { # $1 = mode: micro (high-concurrency Olla-bound) | stream (realistic, backend-bound) + [ "$PROFILE" = 1 ] || return 0 + local mode="$1" bin="build/regression/olla${EXE}" secs workers + case "$mode" in + micro) secs=20; workers=8 ;; # Olla-bound endpoints sustain high concurrency + stream) secs=25; workers=3 ;; # streaming is backend-bound; keep concurrency low + *) return 0 ;; + esac + local cpu="${LOGDIR}/cpu-${LEG}-${mode}.pb.gz" allocs="${LOGDIR}/allocs-${LEG}-${mode}.pb.gz" + local allocs_base="${LOGDIR}/allocs-${LEG}-${mode}-base.pb.gz" + # Baseline alloc snapshot BEFORE the burst. The alloc_objects profile is + # cumulative-since-boot, so on its own it is swamped by one-shot startup work + # (YAML config parse). Diffing after-vs-base with `go tool pprof -base` isolates + # the allocations made *during* the burst - the true per-request signal. + curl -sf --max-time 10 "${PPROF}/allocs" -o "$allocs_base" 2>/dev/null + # No -f: on a non-2xx we still want the body (error text) on disk so an empty + # capture is diagnosable; the HTTP status is saved alongside. (Go's CPU profiler + # can be flaky on Windows under heavy syscall load - allocs/goroutine are not.) + curl -s --max-time $((secs + 15)) -w '%{http_code}' "${PPROF}/profile?seconds=${secs}" \ + -o "$cpu" > "${LOGDIR}/cpu-${LEG}-${mode}.status" 2>/dev/null & + local prof_pid=$! + local deadline=$(( $(date +%s) + secs )) i=0 worker_pids="" + while [ "$i" -lt "$workers" ]; do + ( while [ "$(date +%s)" -lt "$deadline" ]; do perf_request "$mode"; done ) & + worker_pids="$worker_pids $!" + i=$((i + 1)) + done + wait "$prof_pid" 2>/dev/null + curl -sf --max-time 10 "${PPROF}/allocs" -o "$allocs" 2>/dev/null + # Wait only for the worker subshells - bare 'wait' would also wait for the Olla + # background process (started in boot_olla) and hang until the server exits. + # shellcheck disable=SC2086 + wait $worker_pids 2>/dev/null || true + go tool pprof -top -cum -nodecount=12 "$bin" "$cpu" > "${LOGDIR}/cpu-${LEG}-${mode}-top.txt" 2>/dev/null + # Cumulative (since-boot) and the burst-only delta (after vs base). The delta is + # the per-request signal; the cumulative is kept for reference. + go tool pprof -top -alloc_objects -nodecount=12 "$bin" "$allocs" > "${LOGDIR}/allocs-${LEG}-${mode}-cumulative-top.txt" 2>/dev/null + go tool pprof -top -alloc_objects -base "$allocs_base" -nodecount=12 "$bin" "$allocs" > "${LOGDIR}/allocs-${LEG}-${mode}-delta-top.txt" 2>/dev/null + # Report CPU and allocs independently - allocs is the reliable signal on Windows. + if [ -s "${LOGDIR}/cpu-${LEG}-${mode}-top.txt" ]; then + assert_pass "[${LEG}/${mode}] PERF CPU profile captured" + echo "INFO: [${LEG}/${mode}] top CPU (cum):"; sed -n '1,9p' "${LOGDIR}/cpu-${LEG}-${mode}-top.txt" + else + assert_warn "[${LEG}/${mode}] PERF CPU profile empty (http=$(tr -d '\r\n' < "${LOGDIR}/cpu-${LEG}-${mode}.status" 2>/dev/null); CPU profiling can be flaky on Windows)" + fi + if [ -s "${LOGDIR}/allocs-${LEG}-${mode}-delta-top.txt" ]; then + assert_pass "[${LEG}/${mode}] PERF allocs delta captured (burst-only, startup excluded)" + echo "INFO: [${LEG}/${mode}] top allocators during burst:"; sed -n '1,9p' "${LOGDIR}/allocs-${LEG}-${mode}-delta-top.txt" + else + assert_warn "[${LEG}/${mode}] PERF allocs delta empty (allocs endpoint unreachable or go tool pprof failed)" + fi +} + +perf_leg_end() { # leak/heap check after the full leg (incl. Tier-B), vs the boot baseline + [ "$PROFILE" = 1 ] || return 0 + sleep 5 # settle so transient request goroutines unwind before sampling + local gend=$(pprof_goroutines) hend=$(pprof_heap_field HeapAlloc) gc=$(pprof_heap_field NumGC) + if [ -n "$GBASE" ] && [ -n "$gend" ]; then + local gd=$((gend - GBASE)) + if [ "$gd" -gt 30 ]; then + assert_warn "[${LEG}] PERF goroutines grew ${GBASE}->${gend} (+${gd}) after settle - possible leak" + else + assert_pass "[${LEG}] PERF goroutines stable ${GBASE}->${gend} (+${gd})" + fi + fi + echo "INFO: [${LEG}] PERF heap ${HBASE:-?}->${hend:-?} bytes, NumGC=${gc:-?}" +} +``` + +### Drive both legs + +```bash +for LEG in passthrough translation; do + PT=$([ "$LEG" = passthrough ] && echo true || echo false) + CFG=$(render_config "$PT") + boot_olla "$CFG" || { stop_olla; continue; } + GBASE=""; HBASE="" + if [ "$PROFILE" = 1 ]; then GBASE=$(pprof_goroutines); HBASE=$(pprof_heap_field HeapAlloc); fi + tier_a + perf_burst micro # no-op unless --profile; isolates Olla's CPU (no backend wait) + perf_burst stream # no-op unless --profile; realistic streaming translation/proxy path + tier_b + perf_leg_end # no-op unless --profile; goroutine/heap leak check vs baseline + stop_olla +done +``` + +(Set `LEG` and `PT` as shown; the orchestrator runs Olla in background and +records `OLLA_PID` so `stop_olla`/cleanup can kill it. Restarting between legs +is what flips passthrough<->translation against the one Ollama.) + +## Phase 5 - WARN summary + report + verdict + +Print the WARN summary block to the console first: + +```bash +echo "===== WARN SUMMARY =====" +printf '%b\n' "$WARNS" +echo "========================" +``` + +Write `$REPORT`: + +```markdown +# Claude Code E2E through Olla - +- Olla commit: Branch: +- Ollama: Model: Context: +- Legs: passthrough, translation +- Verdict: PASS | FAIL +- Totals:

P / F / W + +## Tier-A protocol (gate) - per leg +| Check | Passthrough | Translation | +|---|---|---| +| A1 models | ... | ... | +| A2 non-stream | ... | ... | +| A2b X-Olla-Mode header | ... | ... | +| A3 SSE order | ... | ... | +| A4 tool_use | ... | ... | +| A5 count_tokens | ... | ... | +| A6 routing stats delta | ... | ... | + +## Tier-B agent (non-gating) - per leg +| Check | Passthrough | Translation | +|---|---|---| +| B1 claude success | ... | ... | +| B2 tests green | ... | ... | +| B3 summary.go changed | ... | ... | +| B4 TotalSize test added | ... | ... | + +## Performance (--profile only, informational) - per leg +| Metric | Passthrough | Translation | +|---|---|---| +| goroutines baseline -> end (delta) | ... | ... | +| heap baseline -> end | ... | ... | +| NumGC over leg | ... | ... | +| top CPU (cum) path | ... | ... | +| top allocator | ... | ... | + +Raw profiles per leg and mode in the log dir: `cpu--.pb.gz` +and, for allocations, a `-base` snapshot plus the final `allocs--.pb.gz` +(`go tool pprof build/regression/olla ` to explore). Headline text: +`cpu-*-top.txt`, `allocs-*-delta-top.txt` (**burst-only, startup excluded - the +per-request signal**) and `allocs-*-cumulative-top.txt` (since-boot, reference). +Compare passthrough vs translation for the translation path's added overhead, and +micro vs stream to separate per-request middleware cost from the real streaming +allocation profile. + +## Failures + + +## Warnings / Notes + +``` + +`last-runs.md` line is written by the cleanup trap. + +**Verdict rule:** any **Tier-A** FAIL, the fitness gate, the build, the clone +baseline, or a leg that never booted -> overall **FAIL**. Tier-B is reported +and shapes the headline but never fails the gate (the model, not Olla, is the +variable there). Finish by telling the user the verdict, the totals, the three +most important findings in plain sentences, and the report path. + +## Verification protocol (per the skills spec - do before trusting a green) + +1. Dispatch a **separate Sonnet agent** for a read-only audit against the + golden-standard checklist (OS portability, no absolute paths, jq-only JSON, + trap cleanup, status categorisation). Fix everything it flags. +2. Dispatch a **second Sonnet agent** to execute end-to-end against a real + Ollama and fix the runtime issues it finds. +3. Iterate until a cold re-run yields the same verdict. diff --git a/.claude/skills/claude-code-e2e/config.regression.yaml.tmpl b/.claude/skills/claude-code-e2e/config.regression.yaml.tmpl new file mode 100644 index 00000000..5df5ad2b --- /dev/null +++ b/.claude/skills/claude-code-e2e/config.regression.yaml.tmpl @@ -0,0 +1,108 @@ +# Olla config template for the /claude-code-e2e skill. +# +# Rendered at runtime into test/regression/tmp/ with these placeholders +# substituted (the skill does the substitution, never edit the rendered copy): +# __OLLA_PORT__ - loopback port Olla listens on +# __OLLAMA_URL__ - the real Ollama endpoint (--ollama-endpoint) +# __PASSTHROUGH__ - true (native Anthropic passthrough) | false (forced translation) +# +# One real Ollama endpoint, one Olla instance. The skill renders this twice - +# once with __PASSTHROUGH__=true, once =false - and restarts Olla between the +# two legs so both the passthrough and the OpenAI<->Anthropic translation paths +# are exercised against the same backend and model. + +server: + host: "127.0.0.1" + port: __OLLA_PORT__ + read_timeout: 300s # gemma-class models on CPU can be slow to first token + read_header_timeout: 10s + write_timeout: 0s # 0 = no write deadline, required for SSE streaming + shutdown_timeout: 5s + request_logging: true + request_limits: + # Claude Code system prompts + tool schemas are large; keep this generous. + max_body_size: 10485760 + max_header_size: 524288 + rate_limits: + # Wide open: this harness is a correctness gate, not a rate-limit test. + global_requests_per_minute: 100000 + per_ip_requests_per_minute: 50000 + health_requests_per_minute: 100000 + burst_size: 1000 + cleanup_interval: 5m + trust_proxy_headers: false + trusted_proxy_cidrs: ["127.0.0.0/8"] + +proxy: + engine: "olla" + profile: "auto" + load_balancer: "priority" + stream_buffer_size: 8192 + connection_timeout: 10s + response_timeout: 300s # long generations through a local model + read_timeout: 300s + response_header_timeout: 30s + retry: + enabled: true + on_connection_failure: true + max_attempts: 0 + sticky_sessions: + # Off on purpose: Claude Code resends full history each turn, so pinning + # adds a variable without adding coverage for this gate. + enabled: false + +discovery: + type: "static" + refresh_interval: 10s + health_check: + initial_delay: 1s + static: + endpoints: + - url: "__OLLAMA_URL__" + name: "ollama-regression" + type: "ollama" + priority: 100 + model_url: "/api/tags" + health_check_url: "/" + check_interval: 5s + check_timeout: 3s + + model_discovery: + # Fast interval so the model is discoverable within the skill's 60s + # readiness window on a cold boot (matches check_interval). + enabled: true + interval: 5s + timeout: 10s + concurrent_workers: 2 + retry_attempts: 2 + retry_backoff: 1s + +model_registry: + type: "memory" + enable_unifier: true + unification: + enabled: true + stale_threshold: 24h + cleanup_interval: 10m + routing_strategy: + type: "optimistic" + options: + fallback_behavior: "all" + discovery_timeout: 5s + discovery_refresh_on_miss: false + +translators: + anthropic: + enabled: true + passthrough_enabled: __PASSTHROUGH__ + max_message_size: 10485760 + inspector: + enabled: false + +logging: + level: "info" + format: "text" + output: "stdout" + +engineering: + show_nerdstats: false diff --git a/.claude/skills/claude-code-e2e/tasks/totalsize.md b/.claude/skills/claude-code-e2e/tasks/totalsize.md new file mode 100644 index 00000000..dca9d9b6 --- /dev/null +++ b/.claude/skills/claude-code-e2e/tasks/totalsize.md @@ -0,0 +1,46 @@ +# Tier-B task: Summary.TotalSize() + +Fixed feature task run by the child Claude Code instance against a fresh clone +of `thushan/smash`. Chosen because it is one self-contained file plus its test +(`pkg/analysis/summary.go` + `summary_test.go`), depends only on stdlib, and +needs only those two files in context - so it stays inside a small local +model's window while still forcing a real multi-turn tool loop +(read file -> read test -> edit both -> run test). + +## Prompt + +The canonical prompt is the single source of truth in `totalsize.prompt.txt` +(this directory). The skill reads it with `cat` and passes it verbatim to +`claude -p` - do not duplicate the wording here or inline in the skill, so the +two cannot drift. Current content: + +> Add a method `func (t *Summary) TotalSize() uint64` to pkg/analysis/summary.go +> that returns the sum of the Size field of every item the Summary holds. Then +> add a test in pkg/analysis/summary_test.go, following the existing style, that +> verifies the returned total. Run `go test ./pkg/analysis/...` and make sure it +> passes before you finish. + +## Assertion (run by the skill after the child exits, in the clone dir) + +All of these must hold for Tier-B PASS on a leg: + +1. The child reported success - the final `stream-json` event is + `{"type":"result", ... "is_error": false}`. +2. `go test ./pkg/analysis/...` exits 0 in the clone. +3. `pkg/analysis/summary.go` shows as modified in `git status --porcelain` (the + method was actually added, not a no-op run). +4. `pkg/analysis/summary_test.go` shows as modified AND contains `TotalSize` - + otherwise a model could add the method, skip the test, and still pass 2 and 3 + on the strength of the pre-existing suite. + +A leg that fails any of the three is a Tier-B FAIL for that leg, but Tier-B is +non-gating: it is reported and influences the headline, while the run verdict is +set by Tier-A (the protocol gate) and the fitness preflight. Record which of the +three conditions failed so a model-capability failure (condition 1/2) is +distinguishable from an Olla wire failure (which Tier-A would also catch). + +## Reset between legs + +The clone is thrown away and re-cloned (or `git reset --hard && git clean +-xfd`) before the second leg, so the passthrough and translation legs both start +from the identical pinned baseline. diff --git a/.claude/skills/claude-code-e2e/tasks/totalsize.prompt.txt b/.claude/skills/claude-code-e2e/tasks/totalsize.prompt.txt new file mode 100644 index 00000000..29eb7666 --- /dev/null +++ b/.claude/skills/claude-code-e2e/tasks/totalsize.prompt.txt @@ -0,0 +1 @@ +Add a method `func (t *Summary) TotalSize() uint64` to pkg/analysis/summary.go that returns the sum of the Size field of every item the Summary holds. Then add a test in pkg/analysis/summary_test.go, following the existing style in that file, that verifies the returned total. Run `go test ./pkg/analysis/...` and make sure it passes before you finish. diff --git a/.claude/skills/olla-validate/areas/anthropic.md b/.claude/skills/olla-validate/areas/anthropic.md index 3dd8844a..0680b4c2 100644 --- a/.claude/skills/olla-validate/areas/anthropic.md +++ b/.claude/skills/olla-validate/areas/anthropic.md @@ -30,8 +30,14 @@ Headers: `Content-Type: application/json`, `x-api-key: validate`, valid order: `message_start` → `content_block_start` → `content_block_delta`(×N) → `content_block_stop` → `message_delta` → `message_stop`; deltas assemble to non-empty text. -4. `GET /olla/anthropic/v1/models` → 200, Anthropic-format model list - (non-empty `data[]`). +4. `GET /olla/anthropic/v1/models` → 200, real Anthropic wire format: + - Top-level envelope has `has_more` (boolean) and `first_id`/`last_id`. + - `data[]` is non-empty. + - Each entry has `type: "model"` (not `"chat"` — the official Anthropic SDK + uses this as a deserialise discriminator; a wrong value breaks client + compatibility), a non-empty `display_name`, and `created_at` as an ISO + 8601 string. + FAIL if any of the above are missing or wrong. 5. `POST /olla/anthropic/v1/messages/count_tokens` with the standard body → 200 with `input_tokens > 0`. 6. Invalid body (`{"model":"test-model"}` - no messages/max_tokens) → 4xx diff --git a/.claude/skills/olla-validate/areas/limits-failures.md b/.claude/skills/olla-validate/areas/limits-failures.md index 3a43a3a7..672443e4 100644 --- a/.claude/skills/olla-validate/areas/limits-failures.md +++ b/.claude/skills/olla-validate/areas/limits-failures.md @@ -22,19 +22,24 @@ long" on Git Bash/Windows. 1. Baseline: `POST http://127.0.0.1:41142/olla/proxy/v1/chat/completions` with the standard small body → 200. 2. Body cap, declared length: same route with a ~300KB body (pad one message - with filler) → rejected. Current contract: the security chain rejects a - declared oversized Content-Length with a blanket **403** ("Security - validation failed", `internal/app/handlers/application.go`); 413 is also - acceptable. Any 2xx or 5xx = FAIL. + with filler) → rejected with **413** (Content Too Large, RFC 9110 §15.5.14). + A 403 is a FAIL (it was the pre-fix bug). Any 2xx or 5xx = FAIL. 3. Just-under cap (~200KB) → 200. -4. 429 rate limit: fire 45 rapid sequential requests → at least one 429; - record at which request it first appears (expect after roughly - burst+window allowance) and whether a Retry-After or rate-limit header is - present (note, don't fail on header absence). +4. 429 rate limit: fire 45 rapid sequential requests → at least one 429 + (RFC 6585 §4); record at which request it first appears (expect after roughly + burst+window allowance). The 429 response must carry at least one of + `Retry-After` or an `X-RateLimit-*` header — FAIL if neither is present. + A 403 in place of 429 is a FAIL (it was the pre-fix masking bug). 5. Health exemption: while rate-limited, `GET /internal/health` on 41142 still 200 (health has its own generous limit). -6. Malformed JSON (`{"model":`) → 4xx, not 5xx. -7. Empty body POST → 4xx, not 5xx. +6. Malformed JSON (`{"model":`) → must not 5xx and must not hang; record the + observed status (2xx or 4xx are both acceptable). Olla is a transparent + proxy and delegates body validation to the backend by design — it + opportunistically extracts the model field and, on failure, forwards to a + fallback backend. Only a 5xx or a hang is a FAIL. +7. Empty body POST → must not 5xx and must not hang; record the observed + status (2xx or 4xx are both acceptable, same delegation rationale as + check 6). Only a 5xx or a hang is a FAIL. 8. Wrong method: `DELETE /olla/proxy/v1/chat/completions` → 404/405, no 5xx. ### Nightly additions diff --git a/.claude/skills/olla-validate/areas/openai-api.md b/.claude/skills/olla-validate/areas/openai-api.md index 88419277..c74ffcd2 100644 --- a/.claude/skills/olla-validate/areas/openai-api.md +++ b/.claude/skills/olla-validate/areas/openai-api.md @@ -24,8 +24,11 @@ a check without mutating mock state, mark it skip with a note). 5. Ollama native non-stream: `POST /olla/ollama/api/chat` (`"stream":false`) → 200, `done: true`, `message.content` non-empty, `prompt_eval_count`/`eval_count` > 0. -6. Malformed body to `/olla/proxy/v1/chat/completions` (`{"model":`) → 4xx - (FAIL if 5xx or hang). +6. Malformed body to `/olla/proxy/v1/chat/completions` (`{"model":`) → must + not 5xx and must not hang; record the observed status (2xx or 4xx are both + acceptable). Olla is a transparent proxy and delegates body validation to + the backend by design — it opportunistically extracts the model field and, + on failure, forwards to a fallback backend. Only a 5xx or a hang is a FAIL. ## Nightly additions diff --git a/.gitignore b/.gitignore index 88bb38fd..7420d4d0 100644 --- a/.gitignore +++ b/.gitignore @@ -41,3 +41,6 @@ config.yaml *.local.json *.log /test/reports/ + +# /claude-code-e2e throwaway clones + rendered configs +test/regression/tmp/ diff --git a/docs/content/api-reference/anthropic.md b/docs/content/api-reference/anthropic.md index 02034b44..2958e457 100644 --- a/docs/content/api-reference/anthropic.md +++ b/docs/content/api-reference/anthropic.md @@ -377,6 +377,52 @@ curl -X POST http://localhost:40114/olla/anthropic/v1/messages \ **Note**: Tool use requires a model that supports function calling. Not all local models support this feature. +### Example: Reasoning (Thinking Blocks) + +Reasoning models (DeepSeek-R1, Qwen3, gpt-oss) emit a chain-of-thought before their answer. Olla surfaces it as a `thinking` block ahead of the text. Passthrough forwards the native block as-is; translation builds one from the backend's `reasoning`/`reasoning_content` field. + +**Non-streaming response**: + +```json +{ + "id": "msg_reason123", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "The user wants the capital of France. That is Paris." + }, + { + "type": "text", + "text": "The capital of France is Paris." + } + ], + "model": "deepseek-r1:latest", + "stop_reason": "end_turn", + "usage": { "input_tokens": 12, "output_tokens": 34 } +} +``` + +**Streaming**: the thinking block streams first, before the text block: + +``` +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"The user wants"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} +``` + +**Notes**: + +- Local models don't produce Anthropic's `signature`, so thinking blocks come through without one. +- Olla reads both field names: `reasoning` (Ollama, LM Studio, Lemonade) and `reasoning_content` (vLLM, SGLang, DeepSeek). +- Reasoning tokens count against `max_tokens`. Set it too low and the model spends the lot thinking, returning `stop_reason: "max_tokens"` with little or no text, so give reasoning models headroom. + ### Example: Vision (Image Input) **Request**: @@ -587,7 +633,7 @@ Streaming uses Server-Sent Events (SSE) with typed events. |-------|-------------|--------------| | `message_start` | Initial message metadata | `{"type":"message_start","message":{...}}` | | `content_block_start` | Start of text or tool_use block | `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` | -| `content_block_delta` | Text chunks (`text_delta`) or tool JSON chunks (`input_json_delta`) | `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"..."}}` | +| `content_block_delta` | Text chunks (`text_delta`), reasoning chunks (`thinking_delta`), or tool JSON chunks (`input_json_delta`) | `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"..."}}` | | `content_block_stop` | End of content block | `{"type":"content_block_stop","index":0}` | | `message_delta` | Stop reason and final usage statistics | `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":9}}` | | `message_stop` | End of stream | `{"type":"message_stop"}` | @@ -712,6 +758,7 @@ Errors follow Anthropic API format: - System messages (string or content blocks) - Tool use (definitions, tool_choice, tool_use, tool_result) - Tool streaming with `input_json_delta` events +- **Reasoning output** surfaced as `thinking` blocks (passthrough keeps the native block; translation builds one from `reasoning`/`reasoning_content`, streamed as `thinking_delta`) - Token counting via `/count_tokens` endpoint - Stop sequences - Temperature, top_p, top_k parameters @@ -733,7 +780,7 @@ Errors follow Anthropic API format: ### ❌ Not Supported -- **Extended Thinking**: Field accepted but not processed +- **Extended Thinking control**: the request-side `thinking`/`budget_tokens` parameter is accepted but ignored. (Reasoning *output* is still surfaced as `thinking` blocks, see above.) - **Prompt Caching**: Not implemented - **Batches API**: Not implemented - **Message Editing**: Not supported diff --git a/docs/content/concepts/api-translation.md b/docs/content/concepts/api-translation.md index 7396d84a..4de87034 100644 --- a/docs/content/concepts/api-translation.md +++ b/docs/content/concepts/api-translation.md @@ -510,7 +510,7 @@ Not all features translate perfectly: **Anthropic → OpenAI Limitations**: -- Extended thinking: Not supported (Anthropic-specific feature) +- Extended thinking control: the request-side `thinking`/`budget_tokens` is ignored. Reasoning *output* is preserved though, surfaced as `thinking` blocks from the backend's `reasoning`/`reasoning_content` field. - Prompt caching: Not supported (Anthropic-specific feature) - Some advanced parameters may not have OpenAI equivalents diff --git a/docs/content/integrations/api-translation/anthropic.md b/docs/content/integrations/api-translation/anthropic.md index fa29ff50..d23a0440 100644 --- a/docs/content/integrations/api-translation/anthropic.md +++ b/docs/content/integrations/api-translation/anthropic.md @@ -694,7 +694,7 @@ See [Crush CLI Integration](../frontend/crush-cli.md) for complete setup. ### Not Supported -- ❌ **Extended Thinking** - Advanced reasoning mode +- ❌ **Extended Thinking control** - request-side `thinking`/`budget_tokens` is ignored. (Reasoning *output* is supported, surfaced as `thinking` blocks in both modes.) - ❌ **Prompt Caching** - Response caching - ❌ **Batches API** - Batch processing - ❌ **Usage Tracking** - Account-level usage tracking diff --git a/internal/adapter/health/checker.go b/internal/adapter/health/checker.go index c122bb58..787461d7 100644 --- a/internal/adapter/health/checker.go +++ b/internal/adapter/health/checker.go @@ -128,14 +128,15 @@ func (c *HTTPHealthChecker) StopChecking(ctx context.Context) error { } func (c *HTTPHealthChecker) healthCheckLoop(ctx context.Context) { + // Function-level safety net: if something truly unrecoverable escapes the + // per-tick inner recover, mark isRunning false so StartChecking can restart. defer func() { if c.ticker != nil { c.ticker.Stop() } - // Panic recovery for health check loop if r := recover(); r != nil { - c.logger.Error("Health check loop panic recovered", "panic", r) - // Could restart the loop here if needed + c.logger.Error("Health check loop exited unexpectedly", "panic", r) + c.isRunning.Store(false) } }() @@ -148,10 +149,23 @@ func (c *HTTPHealthChecker) healthCheckLoop(ctx context.Context) { c.logger.Debug("Health check loop stopping due to stop signal") return case <-c.ticker.C: - // Use a separate context for health checks to avoid cancelling mid-check - checkCtx, cancel := context.WithTimeout(context.Background(), DefaultHealthCheckInterval/2) - c.performHealthChecks(checkCtx) - cancel() + // Wrap per-tick work in its own recover so a panic in + // performHealthChecks does not kill the goroutine - the loop + // continues and the next tick fires normally. + func() { + defer func() { + if r := recover(); r != nil { + c.logger.Error("Health check tick panic recovered, loop continues", + "panic", r) + } + }() + // Derive from the loop's own context so in-flight checks are + // cancelled when the checker shuts down, rather than running + // to their full timeout against an already-stopped service. + checkCtx, cancel := context.WithTimeout(ctx, DefaultHealthCheckInterval/2) + defer cancel() + c.performHealthChecks(checkCtx) + }() } } } @@ -305,7 +319,7 @@ func (c *HTTPHealthChecker) checkEndpoint(ctx context.Context, endpoint *domain. // Trigger unhealthy callback only when an endpoint becomes non-routable from a // routable state. Busy and Warming are still routable, so a Healthy→Busy transition - // must not evict sticky sessions — that would defeat KV-cache affinity entirely. + // must not evict sticky sessions - that would defeat KV-cache affinity entirely. // Unknown→Unhealthy is intentionally excluded: nothing could have been pinned to an // endpoint that was never routable, so there is nothing to purge. if statusChanged && !newStatus.IsRoutable() && oldStatus.IsRoutable() { diff --git a/internal/adapter/health/checker_test.go b/internal/adapter/health/checker_test.go index ed9f5c9b..e5cded41 100644 --- a/internal/adapter/health/checker_test.go +++ b/internal/adapter/health/checker_test.go @@ -734,7 +734,7 @@ func TestStopChecking_DoubleInvoke(t *testing.T) { t.Fatalf("StartChecking: %v", err) } - // Two concurrent stops — neither should panic. + // Two concurrent stops - neither should panic. var wg sync.WaitGroup wg.Add(2) for range 2 { @@ -745,3 +745,149 @@ func TestStopChecking_DoubleInvoke(t *testing.T) { } wg.Wait() } + +// panicRepository panics on GetAll after a configurable number of successful calls, +// then returns normally - used to verify the healthCheckLoop survives a tick panic. +type panicRepository struct { + *mockRepository + callsUntilPanic int + calls int + mu sync.Mutex +} + +func (p *panicRepository) GetAll(ctx context.Context) ([]*domain.Endpoint, error) { + p.mu.Lock() + p.calls++ + calls := p.calls + p.mu.Unlock() + + if calls == p.callsUntilPanic { + panic("injected panic in GetAll") + } + return p.mockRepository.GetAll(ctx) +} + +// TestHealthCheckLoop_SurvivesPanic verifies that a panic inside performHealthChecks +// does not kill the healthCheckLoop goroutine. The loop must continue firing on +// subsequent ticks and the isRunning flag must remain true after the panic. +func TestHealthCheckLoop_SurvivesPanic(t *testing.T) { + t.Parallel() + + loggerCfg := &logger.Config{Level: "error", Theme: "default"} + log, cleanup, _ := logger.New(loggerCfg) + defer cleanup() + styledLogger := logger.NewPlainStyledLogger(log) + + panicRepo := &panicRepository{ + mockRepository: newMockRepository(), + callsUntilPanic: 1, // panic on the first tick + } + + // Use a very short ticker interval so the test doesn't have to wait long. + checker := NewHTTPHealthChecker(panicRepo, styledLogger, &mockHTTPClient{statusCode: 200}) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + checker.isRunning.Store(true) + checker.ticker = time.NewTicker(20 * time.Millisecond) + go checker.healthCheckLoop(ctx) + + // Wait long enough for two ticks (the first panics, the second must succeed). + time.Sleep(120 * time.Millisecond) + + // isRunning must still be true - the loop survived the panic. + if !checker.isRunning.Load() { + t.Error("isRunning is false after tick panic; loop likely died") + } + + // The second GetAll call (tick 2) must have happened, proving the loop continued. + panicRepo.mu.Lock() + calls := panicRepo.calls + panicRepo.mu.Unlock() + + if calls < 2 { + t.Errorf("GetAll called %d times; expected >= 2 (loop must have continued past the panic)", calls) + } + + _ = checker.StopChecking(ctx) +} + +// blockingHTTPClient blocks in Do until the request context is cancelled, then +// records that cancellation was observed. Used to verify that in-flight checks +// are cancelled when the loop context is cancelled rather than running to their +// full timeout. +type blockingHTTPClient struct { + cancelled chan struct{} // closed when a Do call observes ctx cancellation +} + +func (b *blockingHTTPClient) Do(req *http.Request) (*http.Response, error) { + <-req.Context().Done() + select { + case <-b.cancelled: + default: + close(b.cancelled) + } + return nil, req.Context().Err() +} + +// TestHealthCheckLoop_CancelledContextAbortsInFlightChecks verifies that when the +// loop context is cancelled, any in-flight health check context is also cancelled +// rather than running to its full per-tick timeout. Before the fix, checkCtx was +// derived from context.Background() so cancelling the loop ctx had no effect on +// the check in progress. +func TestHealthCheckLoop_CancelledContextAbortsInFlightChecks(t *testing.T) { + if testing.Short() { + t.Skip("skipping shutdown-cancellation test in short mode") + } + t.Parallel() + + loggerCfg := &logger.Config{Level: "error", Theme: "default"} + log, cleanup, _ := logger.New(loggerCfg) + defer cleanup() + styledLogger := logger.NewPlainStyledLogger(log) + + repo := newMockRepository() + + testURL, _ := url.Parse("http://localhost:19999") + healthURL, _ := url.Parse("http://localhost:19999/health") + ep := &domain.Endpoint{ + Name: "blocking-ep", + URL: testURL, + HealthCheckURL: healthURL, + CheckTimeout: 5 * time.Second, + URLString: testURL.String(), + } + repo.mu.Lock() + repo.endpoints[testURL.String()] = ep + repo.mu.Unlock() + + blocking := &blockingHTTPClient{cancelled: make(chan struct{})} + checker := NewHTTPHealthChecker(repo, styledLogger, blocking) + + ctx, cancel := context.WithCancel(context.Background()) + + // Drive a single performHealthChecks tick directly from a goroutine, then + // cancel the ctx immediately. We skip the ticker so the test is deterministic. + started := make(chan struct{}) + go func() { + close(started) + // Per-tick timeout must also be short so cancellation is observable. + checkCtx, checkCancel := context.WithTimeout(ctx, 5*time.Second) + defer checkCancel() + checker.performHealthChecks(checkCtx) + }() + + // Wait for the goroutine to start, then cancel. + <-started + cancel() + + // The in-flight Do call should observe the cancellation well within the + // per-tick timeout (5 s). Allow 500 ms as a generous but bounded window. + select { + case <-blocking.cancelled: + // pass - in-flight check was cancelled + case <-time.After(500 * time.Millisecond): + t.Fatal("in-flight health check was not cancelled within 500ms after loop context was cancelled") + } +} diff --git a/internal/adapter/inspector/body_inspector.go b/internal/adapter/inspector/body_inspector.go index d06146b3..48221744 100644 --- a/internal/adapter/inspector/body_inspector.go +++ b/internal/adapter/inspector/body_inspector.go @@ -27,7 +27,7 @@ const ( // when "model" appears after large fields (e.g. a multi-megabyte "messages" array); // 8 MB accommodates typical vision payloads while staying well within practical limits. // Requests where "model" is absent or appears beyond this ceiling will not be routed by - // model name — capability-based or default routing applies instead. + // model name - capability-based or default routing applies instead. modelScanSize = 8 * 1024 * 1024 ) @@ -80,7 +80,7 @@ func (bi *BodyInspector) Inspect(ctx context.Context, r *http.Request, profile * // them. This handles the case where "messages" (with large base64 images) precedes "model". if r.ContentLength > bi.maxBodySize { // captured is a locally-allocated bytes.Buffer (not from a sync.Pool), so - // captured.Bytes() is safe to use directly without an extra copy — unlike the + // captured.Bytes() is safe to use directly without an extra copy - unlike the // small-body path below which must copy buffer.Bytes() to avoid pool aliasing // once the deferred Reset()/Put() runs. captured := &bytes.Buffer{} @@ -108,7 +108,10 @@ func (bi *BodyInspector) Inspect(ctx context.Context, r *http.Request, profile * return nil } - buffer := bi.bufferPool.Get() + buffer, err := bi.bufferPool.Get() + if err != nil { + return fmt.Errorf("body inspector: buffer pool exhausted: %w", err) + } defer func() { buffer.Reset() bi.bufferPool.Put(buffer) @@ -164,7 +167,7 @@ func (bi *BodyInspector) extractModelName(body []byte) string { return "" } - // Fast path: complete JSON — unmarshal directly. + // Fast path: complete JSON - unmarshal directly. var req modelRequest if err := json.Unmarshal(body, &req); err == nil && req.Model != "" { return bi.normalizeModelName(req.Model) @@ -230,7 +233,7 @@ func extractTopLevelModelFieldFromReader(r io.Reader) string { key, ok := keyTok.(string) if !ok { - // Malformed JSON — key position must be a string. + // Malformed JSON - key position must be a string. return "" } @@ -244,14 +247,14 @@ func extractTopLevelModelFieldFromReader(r io.Reader) string { if err := json.Unmarshal(val, &modelStr); err == nil && modelStr != "" { return modelStr } - // Value wasn't a string — keep scanning. + // Value wasn't a string - keep scanning. continue } // Skip the value for this key using token-level iteration so we never buffer // a large nested value (e.g. a messages array containing base64 images). if err := skipValueTokens(dec); err != nil { - // Truncated JSON or I/O error — stop scanning. + // Truncated JSON or I/O error - stop scanning. return "" } } @@ -269,11 +272,11 @@ func skipValueTokens(dec *json.Decoder) error { } delim, ok := tok.(json.Delim) if !ok { - // Scalar value (string, number, bool, null) — already consumed. + // Scalar value (string, number, bool, null) - already consumed. return nil } if delim != '{' && delim != '[' { - // Unexpected closing delimiter at value position — treat as error. + // Unexpected closing delimiter at value position - treat as error. return fmt.Errorf("unexpected closing delimiter %v", delim) } // Track open/close depth to handle arbitrarily nested structures. diff --git a/internal/adapter/proxy/benchmark_refactor_test.go b/internal/adapter/proxy/benchmark_refactor_test.go index 2535104a..7c00aa7a 100644 --- a/internal/adapter/proxy/benchmark_refactor_test.go +++ b/internal/adapter/proxy/benchmark_refactor_test.go @@ -318,7 +318,10 @@ func BenchmarkPoolGetPut(b *testing.B) { b.ReportAllocs() for i := range b.N { - obj := pool.Get() + obj, err := pool.Get() + if err != nil { + b.Fatal(err) + } // Use the object obj.data[0] = byte(i) pool.Put(obj) @@ -370,7 +373,10 @@ func BenchmarkOllaObjectPools(b *testing.B) { b.ReportAllocs() for range b.N { - buf := bufferPool.Get() + buf, err := bufferPool.Get() + if err != nil { + b.Fatal(err) + } bufferPool.Put(buf) } }) @@ -394,7 +400,10 @@ func BenchmarkOllaObjectPools(b *testing.B) { b.ReportAllocs() for range b.N { - ctx := reqPool.Get() + ctx, err := reqPool.Get() + if err != nil { + b.Fatal(err) + } reqPool.Put(ctx) } }) diff --git a/internal/adapter/proxy/core/retry.go b/internal/adapter/proxy/core/retry.go index 53776d3f..b2100b56 100644 --- a/internal/adapter/proxy/core/retry.go +++ b/internal/adapter/proxy/core/retry.go @@ -242,12 +242,22 @@ func (h *RetryHandler) buildFinalError(availableEndpoints []*domain.Endpoint, ma return fmt.Errorf("max attempts (%d) reached: %w", maxRetries, lastErr) } -// IsConnectionError identifies transient network errors suitable for retry +// IsConnectionError identifies transient network errors suitable for retry. +// Client-side context cancellations and deadline exceeded errors are NOT +// connection errors — the backend is not at fault and must not be penalised +// or failed-over with an already-expired context. func IsConnectionError(err error) bool { if err == nil { return false } + // Short-circuit before the net.Error check: context.DeadlineExceeded's + // concrete type (*url.Error wrapping it) satisfies net.Error (Timeout()==true), + // which would falsely mark the endpoint unhealthy for a client timeout. + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return false + } + var netErr net.Error if errors.As(err, &netErr) { return true diff --git a/internal/adapter/proxy/core/retry_safety_test.go b/internal/adapter/proxy/core/retry_safety_test.go index cadf3c1e..6213629f 100644 --- a/internal/adapter/proxy/core/retry_safety_test.go +++ b/internal/adapter/proxy/core/retry_safety_test.go @@ -58,6 +58,65 @@ var _ net.Error = (*connectionResetError)(nil) // ---- tests ------------------------------------------------------------------ +// TestIsConnectionError verifies the context-cancellation fix: client-side timeouts +// and cancellations must NOT be classified as connection errors — the endpoint is +// innocent and must not be penalised or failed-over on an already-expired context. +func TestIsConnectionError(t *testing.T) { + t.Parallel() + + // A genuine net.Error that is NOT a context error. + type urlError struct{ msg string } + _ = urlError{} + + tests := []struct { + name string + err error + want bool + }{ + { + name: "nil error", + err: nil, + want: false, + }, + { + name: "context.DeadlineExceeded", + err: context.DeadlineExceeded, + want: false, + }, + { + name: "context.Canceled", + err: context.Canceled, + want: false, + }, + { + name: "url.Error wrapping context.DeadlineExceeded", + err: &net.OpError{ + Op: "dial", + Err: context.DeadlineExceeded, + }, + want: false, + }, + { + name: "genuine net.Error (connection reset)", + err: &connectionResetError{}, + want: true, + }, + { + name: "generic non-network error", + err: errors.New("something went wrong"), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := IsConnectionError(tt.err) + assert.Equal(t, tt.want, got, "IsConnectionError(%v)", tt.err) + }) + } +} + // TestIsIdempotent confirms the idempotency predicate matches RFC 9110. func TestIsIdempotent(t *testing.T) { t.Parallel() diff --git a/internal/adapter/proxy/olla/service.go b/internal/adapter/proxy/olla/service.go index 86f0a9b3..00cec7e3 100644 --- a/internal/adapter/proxy/olla/service.go +++ b/internal/adapter/proxy/olla/service.go @@ -372,8 +372,11 @@ func (s *Service) UpdateConfig(config ports.ProxyConfiguration) { newConfig.MaxIdleConnsPerHost = ollaConfig.MaxIdleConnsPerHost newConfig.ResponseHeaderTimeout = ollaConfig.ResponseHeaderTimeout newConfig.TLSHandshakeTimeout = ollaConfig.TLSHandshakeTimeout - } else { - // fallback: preserve current Olla-specific settings for non-Olla configs + } else if current != nil { + // fallback: preserve current Olla-specific settings for non-Olla configs. + // Guard against a nil current pointer — only reachable if UpdateConfig is + // called on a zero-value Service (e.g. in tests) before NewService stores + // the initial configuration. newConfig.MaxIdleConns = current.MaxIdleConns newConfig.IdleConnTimeout = current.IdleConnTimeout newConfig.MaxConnsPerHost = current.MaxConnsPerHost @@ -385,11 +388,12 @@ func (s *Service) UpdateConfig(config ports.ProxyConfiguration) { s.configuration.Store(newConfig) } -// cleanupLoop periodically cleans up unused endpoint pools and circuit breakers +// cleanupLoop periodically cleans up unused endpoint pools and circuit breakers. func (s *Service) cleanupLoop() { + // Function-level safety net in case something escapes the per-tick recover. defer func() { if r := recover(); r != nil { - s.Logger.Error("cleanupLoop panic recovered", "panic", r) + s.Logger.Error("cleanupLoop exited unexpectedly", "panic", r) } }() @@ -398,7 +402,17 @@ func (s *Service) cleanupLoop() { case <-s.cleanupStop: return case <-s.cleanupTicker.C: - s.cleanupUnusedResources() + // Wrap per-tick work so a panic in cleanupUnusedResources does not + // kill the goroutine — the loop continues and cleans up next tick. + func() { + defer func() { + if r := recover(); r != nil { + s.Logger.Error("cleanupLoop tick panic recovered, loop continues", + "panic", r) + } + }() + s.cleanupUnusedResources() + }() } } } diff --git a/internal/adapter/proxy/olla/service_concurrency_test.go b/internal/adapter/proxy/olla/service_concurrency_test.go index dce436b2..1a126197 100644 --- a/internal/adapter/proxy/olla/service_concurrency_test.go +++ b/internal/adapter/proxy/olla/service_concurrency_test.go @@ -6,6 +6,7 @@ import ( "time" "github.com/puzpuzpuz/xsync/v4" + "github.com/thushan/olla/internal/adapter/proxy/config" "github.com/thushan/olla/internal/adapter/proxy/core" ) @@ -157,3 +158,37 @@ func TestGetCircuitBreaker_ComputeOnce(t *testing.T) { } } } + +// TestUpdateConfig_NilGuard verifies that calling UpdateConfig on a Service whose +// configuration pointer has never been stored (nil atomic) does not panic. +// This covers the else-branch nil guard added in T0-2: if current is nil we fall +// through without dereferencing it, and the new config is stored correctly. +func TestUpdateConfig_NilGuard(t *testing.T) { + t.Parallel() + + // Construct a Service without going through NewService so the atomic is nil. + s := &Service{ + BaseProxyComponents: &core.BaseProxyComponents{ + Logger: createTestLogger(), + }, + endpointPools: *xsync.NewMap[string, *connectionPool](), + circuitBreakers: *xsync.NewMap[string, *circuitBreaker](), + } + // configuration atomic is zero-value — Load() returns nil. + + // Use a non-*Configuration to trigger the else-branch (which used to dereference nil). + nonOlla := &config.SherpaConfig{} + nonOlla.ReadTimeout = 5 * time.Second + + // Must not panic. + s.UpdateConfig(nonOlla) + + // Config should now be stored. + stored := s.configuration.Load() + if stored == nil { + t.Fatal("expected configuration to be stored after UpdateConfig, got nil") + } + if stored.ReadTimeout != 5*time.Second { + t.Errorf("ReadTimeout: want 5s, got %v", stored.ReadTimeout) + } +} diff --git a/internal/adapter/proxy/olla/service_leak_test.go b/internal/adapter/proxy/olla/service_leak_test.go index 4c28fed2..c64daea6 100644 --- a/internal/adapter/proxy/olla/service_leak_test.go +++ b/internal/adapter/proxy/olla/service_leak_test.go @@ -290,6 +290,50 @@ func (m *mockStatsCollector) GetEndpointStats() map[string]ports.EndpointStats { func (m *mockStatsCollector) GetSecurityStats() ports.SecurityStats { return ports.SecurityStats{} } func (m *mockStatsCollector) GetConnectionStats() map[string]int64 { return nil } +// TestCleanupLoop_SurvivesTick_Panic verifies that a panic inside +// cleanupUnusedResources does not kill the cleanupLoop goroutine. Before the +// per-tick inner recover was added, a single panic would permanently stop cleanup. +func TestCleanupLoop_SurvivesTick_Panic(t *testing.T) { + t.Parallel() + + s := &Service{ + BaseProxyComponents: &core.BaseProxyComponents{ + Logger: createTestLogger(), + }, + endpointPools: *xsync.NewMap[string, *connectionPool](), + circuitBreakers: *xsync.NewMap[string, *circuitBreaker](), + cleanupTicker: time.NewTicker(20 * time.Millisecond), + cleanupStop: make(chan struct{}), + } + s.configuration.Store(&Configuration{}) + + // Insert a pool whose transport is nil — CloseIdleConnections() will panic + // when cleanupUnusedResources tries to close it (staleThreshold check passes + // because lastUsed == 0, which is always in the past). + s.endpointPools.Store("bad-endpoint", &connectionPool{ + transport: nil, // nil dereference in CloseIdleConnections + lastUsed: 0, // triggers cleanup immediately + healthy: 1, + }) + + go s.cleanupLoop() + + // Wait for at least two ticks: first panics, second must still fire. + time.Sleep(100 * time.Millisecond) + + // Liveness probe: an alive loop is parked in its select and receives this + // send (then exits). A dead goroutine never receives it. + select { + case s.cleanupStop <- struct{}{}: + // Expected: loop survived the panic and took the stop signal. + case <-time.After(2 * time.Second): + t.Fatal("cleanupLoop did not receive stop signal; goroutine died after tick panic") + } + + // Shut down cleanly. + s.Cleanup() +} + // TestCleanup_DoubleInvoke verifies that calling Cleanup twice does not panic. // Previously, the second call would close an already-closed channel. func TestCleanup_DoubleInvoke(t *testing.T) { diff --git a/internal/adapter/proxy/olla/service_retry.go b/internal/adapter/proxy/olla/service_retry.go index abd424f3..483f408c 100644 --- a/internal/adapter/proxy/olla/service_retry.go +++ b/internal/adapter/proxy/olla/service_retry.go @@ -138,7 +138,11 @@ func (s *Service) proxyToSingleEndpoint(ctx context.Context, w http.ResponseWrit streamStart := time.Now() stats.FirstDataMs = time.Since(stats.StartTime).Milliseconds() - buffer := s.bufferPool.Get() + buffer, poolErr := s.bufferPool.Get() + if poolErr != nil { + s.RecordFailure(ctx, endpoint, time.Since(stats.StartTime), poolErr) + return fmt.Errorf("olla: stream buffer unavailable: %w", poolErr) + } defer s.bufferPool.Put(buffer) // Separate client and upstream contexts for proper cancellation handling diff --git a/internal/adapter/proxy/sherpa/service_retry.go b/internal/adapter/proxy/sherpa/service_retry.go index 8521ffe4..b10ff112 100644 --- a/internal/adapter/proxy/sherpa/service_retry.go +++ b/internal/adapter/proxy/sherpa/service_retry.go @@ -127,7 +127,11 @@ func (s *Service) proxyToSingleEndpoint(ctx context.Context, w http.ResponseWrit streamStart := time.Now() stats.FirstDataMs = time.Since(stats.StartTime).Milliseconds() - buffer := s.bufferPool.Get() + buffer, poolErr := s.bufferPool.Get() + if poolErr != nil { + s.RecordFailure(ctx, endpoint, time.Since(stats.StartTime), poolErr) + return fmt.Errorf("sherpa: stream buffer unavailable: %w", poolErr) + } defer s.bufferPool.Put(buffer) // Separate client and upstream contexts for proper cancellation handling diff --git a/internal/adapter/stats/collector.go b/internal/adapter/stats/collector.go index 99d7ca93..300e59e7 100644 --- a/internal/adapter/stats/collector.go +++ b/internal/adapter/stats/collector.go @@ -43,6 +43,15 @@ const ( MaxTrackedEndpoints = 50 EndpointTTL = 1 * time.Hour CleanupInterval = 5 * time.Minute + + // MaxUniqueRateLimitedIPs caps the in-memory set of rate-limited client IPs. + // Under an active flood attack the map would otherwise grow without bound. + // At cap, new entries are refused until the periodic age-eviction frees space. + MaxUniqueRateLimitedIPs = 10_000 + + // ipCleanupInterval is how often the age-based eviction scan runs. Keeping + // it off the per-violation hot path avoids an O(N) scan under high load. + ipCleanupInterval = 5 * time.Minute ) type Collector struct { @@ -67,6 +76,7 @@ type Collector struct { rateLimitViolations *xsync.Counter sizeLimitViolations *xsync.Counter lastCleanup int64 + lastIPCleanup int64 // atomic: tracks when the age-eviction scan last ran securityMu sync.RWMutex cleanupMu sync.Mutex @@ -265,16 +275,47 @@ func (c *Collector) RecordSecurityViolation(violation ports.SecurityViolation) { func (c *Collector) recordRateLimitedIP(clientIP string) { now := time.Now().UnixNano() + + // Age-eviction scan runs at most once per ipCleanupInterval, not on every + // violation. The pre-check outside the lock avoids contention when the + // interval has not yet elapsed (the common case under a flood). + lastClean := atomic.LoadInt64(&c.lastIPCleanup) + if now-lastClean >= int64(ipCleanupInterval) { + c.cleanupOldRateLimitedIPs(now) + } + + c.securityMu.Lock() + // Refuse new entries once the hard cap is reached. At cap, the periodic + // eviction will free space for the next cleanup cycle. + if len(c.uniqueRateLimitedIPs) < MaxUniqueRateLimitedIPs { + c.uniqueRateLimitedIPs[clientIP] = now + } + c.securityMu.Unlock() +} + +// cleanupOldRateLimitedIPs removes IPs whose last-seen timestamp is older than +// one hour. Uses a double-checked lock to avoid redundant scans when multiple +// goroutines arrive at the interval boundary concurrently. +func (c *Collector) cleanupOldRateLimitedIPs(now int64) { + c.cleanupMu.Lock() + defer c.cleanupMu.Unlock() + + // Second check after acquiring the cleanup lock. + if now-atomic.LoadInt64(&c.lastIPCleanup) < int64(ipCleanupInterval) { + return + } + cutoff := now - int64(time.Hour) c.securityMu.Lock() - c.uniqueRateLimitedIPs[clientIP] = now for ip, ts := range c.uniqueRateLimitedIPs { if ts < cutoff { delete(c.uniqueRateLimitedIPs, ip) } } c.securityMu.Unlock() + + atomic.StoreInt64(&c.lastIPCleanup, now) } func (c *Collector) updateEndpointStats(endpoint *domain.Endpoint, status string, latencyMs, bytes int64, now int64) { diff --git a/internal/adapter/stats/collector_test.go b/internal/adapter/stats/collector_test.go index 74b12efe..e7840e7f 100644 --- a/internal/adapter/stats/collector_test.go +++ b/internal/adapter/stats/collector_test.go @@ -3,6 +3,7 @@ package stats import ( "net/url" "sync" + "sync/atomic" "testing" "time" @@ -391,3 +392,88 @@ func TestCollector_EmptyStats(t *testing.T) { t.Errorf("Expected 0 rate limit violations, got %d", securityStats.RateLimitViolations) } } + +// TestRecordRateLimitedIP_CapBounded verifies that inserting more than +// MaxUniqueRateLimitedIPs distinct IPs does not grow the map beyond the cap. +// A flood attack must not cause unbounded memory growth. +func TestRecordRateLimitedIP_CapBounded(t *testing.T) { + t.Parallel() + + collector := NewCollector(createTestLogger()) + + // Flood with twice the cap to prove the ceiling holds. + for i := range MaxUniqueRateLimitedIPs * 2 { + collector.recordRateLimitedIP(string(rune('A'+i%26)) + "-" + string(rune('0'+i%10))) + } + + collector.securityMu.RLock() + size := len(collector.uniqueRateLimitedIPs) + collector.securityMu.RUnlock() + + if size > MaxUniqueRateLimitedIPs { + t.Errorf("map size %d exceeds cap %d", size, MaxUniqueRateLimitedIPs) + } +} + +// TestRecordRateLimitedIP_AgeEviction verifies that entries older than one hour +// are removed by the periodic cleanup and GetSecurityStats reports an updated count. +func TestRecordRateLimitedIP_AgeEviction(t *testing.T) { + t.Parallel() + + collector := NewCollector(createTestLogger()) + + // Manually insert an expired entry directly. + collector.securityMu.Lock() + old := time.Now().Add(-2 * time.Hour).UnixNano() + collector.uniqueRateLimitedIPs["expired-ip"] = old + collector.securityMu.Unlock() + + // Force cleanup by resetting lastIPCleanup to zero so the threshold is exceeded. + atomic.StoreInt64(&collector.lastIPCleanup, 0) + collector.cleanupOldRateLimitedIPs(time.Now().UnixNano()) + + stats := collector.GetSecurityStats() + if stats.UniqueRateLimitedIPs != 0 { + t.Errorf("expected 0 unique IPs after eviction, got %d", stats.UniqueRateLimitedIPs) + } +} + +// TestRecordRateLimitedIP_ConcurrentRace exercises concurrent inserts and reads +// under the -race detector to prove no data races on the map. +func TestRecordRateLimitedIP_ConcurrentRace(t *testing.T) { + t.Parallel() + + collector := NewCollector(createTestLogger()) + const goroutines = 20 + const iters = 500 + + var wg sync.WaitGroup + for g := range goroutines { + wg.Add(1) + go func(id int) { + defer wg.Done() + for i := range iters { + collector.RecordSecurityViolation(ports.SecurityViolation{ + ViolationType: constants.ViolationRateLimit, + ClientID: string(rune('A'+id)) + string(rune('0'+i%10)), + }) + } + }(g) + } + // Concurrent readers. + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + for range iters { + _ = collector.GetSecurityStats() + } + }() + } + wg.Wait() + + stats := collector.GetSecurityStats() + if stats.UniqueRateLimitedIPs > MaxUniqueRateLimitedIPs { + t.Errorf("map size %d exceeds cap %d under concurrent load", stats.UniqueRateLimitedIPs, MaxUniqueRateLimitedIPs) + } +} diff --git a/internal/adapter/translator/anthropic/constants.go b/internal/adapter/translator/anthropic/constants.go index 425abd25..e820073e 100644 --- a/internal/adapter/translator/anthropic/constants.go +++ b/internal/adapter/translator/anthropic/constants.go @@ -5,6 +5,7 @@ const ( contentTypeToolUse = "tool_use" contentTypeToolResult = "tool_result" contentTypeImage = "image" + contentTypeThinking = "thinking" ) const ( diff --git a/internal/adapter/translator/anthropic/extended_features_test.go b/internal/adapter/translator/anthropic/extended_features_test.go index a8ce1a8c..daa2e391 100644 --- a/internal/adapter/translator/anthropic/extended_features_test.go +++ b/internal/adapter/translator/anthropic/extended_features_test.go @@ -15,7 +15,7 @@ import ( // TestTransformRequest_WithThinkingField tests that the thinking field is accepted // This verifies the Thinking field is a known field and won't cause "unknown field" errors func TestTransformRequest_WithThinkingField(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) tests := []struct { name string @@ -82,7 +82,7 @@ func TestTransformRequest_WithThinkingField(t *testing.T) { // TestTransformRequest_SystemPromptAsArray tests system prompt with content blocks // Anthropic API supports system prompts as arrays of content blocks func TestTransformRequest_SystemPromptAsArray(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -130,7 +130,7 @@ func TestTransformRequest_SystemPromptAsArray(t *testing.T) { // TestTransformRequest_SystemPromptArrayWithEmptyBlocks tests handling of empty content blocks func TestTransformRequest_SystemPromptArrayWithEmptyBlocks(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -174,7 +174,7 @@ func TestTransformRequest_SystemPromptArrayWithEmptyBlocks(t *testing.T) { // TestTransformRequest_SystemPromptArrayEmpty tests empty system prompt array func TestTransformRequest_SystemPromptArrayEmpty(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -207,7 +207,7 @@ func TestTransformRequest_SystemPromptArrayEmpty(t *testing.T) { // TestTransformRequest_SystemPromptString tests traditional string system prompt still works func TestTransformRequest_SystemPromptString(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -242,7 +242,7 @@ func TestTransformRequest_SystemPromptString(t *testing.T) { // TestTransformRequest_CombinedThinkingAndSystemArray tests both new features together func TestTransformRequest_CombinedThinkingAndSystemArray(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -375,7 +375,7 @@ func TestCountSystemChars(t *testing.T) { // TestTransformRequest_WithContextManagementField is a regression test for GitHub issue #154. // Claude Code v2.1.156+ sends context_management in requests; lenient parsing must not reject it. func TestTransformRequest_WithContextManagementField(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) // Realistic context_management payload as sent by Claude Code. rawJSON := `{ diff --git a/internal/adapter/translator/anthropic/integration_test.go b/internal/adapter/translator/anthropic/integration_test.go index 4156828b..f73e907d 100644 --- a/internal/adapter/translator/anthropic/integration_test.go +++ b/internal/adapter/translator/anthropic/integration_test.go @@ -36,6 +36,17 @@ func createTestConfig() config.AnthropicTranslatorConfig { } } +// mustNewTranslator wraps NewTranslator for test use. NewTranslator now returns +// an error so the no-panic policy is satisfied in production code; tests use +// this helper to keep call-sites terse and fatal on unexpected failures. +func mustNewTranslator(log logger.StyledLogger, cfg config.AnthropicTranslatorConfig) *Translator { + tr, err := NewTranslator(log, cfg) + if err != nil { + panic("mustNewTranslator: " + err.Error()) + } + return tr +} + // setup test http request func createHTTPRequest(t *testing.T, anthropicReq AnthropicRequest) *http.Request { t.Helper() @@ -90,7 +101,7 @@ func simulateBackendResponse(content string, toolCalls []map[string]interface{}, } func TestAnthropicToOpenAIToAnthropic_RoundTrip(t *testing.T) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() t.Run("simple_text_conversation", func(t *testing.T) { @@ -372,7 +383,7 @@ func TestAnthropicToOpenAIToAnthropic_RoundTrip(t *testing.T) { } func TestAnthropicToolCalling_RoundTrip(t *testing.T) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() t.Run("single_tool_definition", func(t *testing.T) { @@ -825,7 +836,7 @@ func TestAnthropicToolCalling_RoundTrip(t *testing.T) { } func TestAnthropicEdgeCases_RoundTrip(t *testing.T) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() t.Run("empty_messages_array", func(t *testing.T) { @@ -978,16 +989,18 @@ func TestAnthropicEdgeCases_RoundTrip(t *testing.T) { transformed, err := translator.TransformRequest(ctx, httpReq) require.NoError(t, err) - // Should create separate messages: user text + tool result + // tool messages must come before the user text so they sit immediately after + // the assistant tool_calls message; OpenAI-compatible backends reject any + // other role between them. messages := transformed.OpenAIRequest["messages"].([]map[string]interface{}) require.Len(t, messages, 2) - assert.Equal(t, "user", messages[0]["role"]) - assert.Equal(t, "Here's the result:", messages[0]["content"]) + assert.Equal(t, "tool", messages[0]["role"]) + assert.Equal(t, "tool_123", messages[0]["tool_call_id"]) + assert.Equal(t, "Result data", messages[0]["content"]) - assert.Equal(t, "tool", messages[1]["role"]) - assert.Equal(t, "tool_123", messages[1]["tool_call_id"]) - assert.Equal(t, "Result data", messages[1]["content"]) + assert.Equal(t, "user", messages[1]["role"]) + assert.Equal(t, "Here's the result:", messages[1]["content"]) }) t.Run("tool_result_with_structured_content", func(t *testing.T) { @@ -1031,7 +1044,7 @@ func TestAnthropicEdgeCases_RoundTrip(t *testing.T) { } func TestAnthropicModelPreservation(t *testing.T) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() models := []string{ @@ -1072,7 +1085,7 @@ func TestAnthropicModelPreservation(t *testing.T) { } func TestAnthropicUsageTracking(t *testing.T) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() anthropicReq := AnthropicRequest{ @@ -1149,7 +1162,7 @@ func TestAnthropicUsageTracking(t *testing.T) { // BenchmarkTransformRequest measures the performance of transforming Anthropic requests to OpenAI format. func BenchmarkTransformRequest(b *testing.B) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() b.Run("simple_text_request", func(b *testing.B) { @@ -1340,7 +1353,7 @@ func BenchmarkTransformRequest(b *testing.B) { // BenchmarkTransformResponse measures the performance of transforming OpenAI responses to Anthropic format. func BenchmarkTransformResponse(b *testing.B) { - translator := NewTranslator(createIntegrationTestLogger(), createTestConfig()) + translator := mustNewTranslator(createIntegrationTestLogger(), createTestConfig()) ctx := context.Background() b.Run("simple_text_response", func(b *testing.B) { diff --git a/internal/adapter/translator/anthropic/passthrough_test.go b/internal/adapter/translator/anthropic/passthrough_test.go index 9e675ff0..6bd9b2f1 100644 --- a/internal/adapter/translator/anthropic/passthrough_test.go +++ b/internal/adapter/translator/anthropic/passthrough_test.go @@ -188,7 +188,7 @@ func TestCanPassthrough(t *testing.T) { PassthroughEnabled: tt.passthroughEnabled, } - translator := NewTranslator(createTestLogger(), cfg) + translator := mustNewTranslator(createTestLogger(), cfg) result := translator.CanPassthrough(tt.endpoints, tt.profileLookup) assert.Equal(t, tt.want, result, tt.description) @@ -471,7 +471,7 @@ func TestPreparePassthrough(t *testing.T) { PassthroughEnabled: true, } - translator := NewTranslator(createTestLogger(), cfg) + translator := mustNewTranslator(createTestLogger(), cfg) // Pre-buffer the body bytes (as the handler does in production) bodyBytes := []byte(tt.requestBody) @@ -515,7 +515,7 @@ func TestPreparePassthrough_OversizedBody(t *testing.T) { PassthroughEnabled: true, } - translator := NewTranslator(createTestLogger(), cfg) + translator := mustNewTranslator(createTestLogger(), cfg) // Build a body that exceeds maxSize oversizedBody := make([]byte, maxSize+1) @@ -551,7 +551,7 @@ func TestPreparePassthrough_WithInspector(t *testing.T) { }, } - translator := NewTranslator(createTestLogger(), cfg) + translator := mustNewTranslator(createTestLogger(), cfg) bodyBytes := []byte(`{ "model": "claude-3-5-sonnet-20241022", @@ -601,7 +601,7 @@ func TestCanPassthrough_Integration(t *testing.T) { PassthroughEnabled: true, } - translator := NewTranslator(createTestLogger(), cfg) + translator := mustNewTranslator(createTestLogger(), cfg) result := translator.CanPassthrough(endpoints, profileLookup) assert.True(t, result, "should support passthrough for vLLM+SGLang deployment") @@ -627,7 +627,7 @@ func TestCanPassthrough_Integration(t *testing.T) { PassthroughEnabled: true, } - translator := NewTranslator(createTestLogger(), cfg) + translator := mustNewTranslator(createTestLogger(), cfg) result := translator.CanPassthrough(endpoints, profileLookup) assert.True(t, result, "should support passthrough when handler pre-filtered to only capable backends") diff --git a/internal/adapter/translator/anthropic/path_translation_test.go b/internal/adapter/translator/anthropic/path_translation_test.go index bb1d4671..0c4babf3 100644 --- a/internal/adapter/translator/anthropic/path_translation_test.go +++ b/internal/adapter/translator/anthropic/path_translation_test.go @@ -15,7 +15,7 @@ import ( // TestPathTranslation verifies that the translator sets the correct target path // This ensures requests are proxied to the correct OpenAI endpoint func TestPathTranslation(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() tests := []struct { @@ -94,7 +94,7 @@ func TestPathTranslation(t *testing.T) { // TestPathTranslationPreservesOtherFields verifies that setting TargetPath doesn't affect other fields func TestPathTranslationPreservesOtherFields(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() anthropicReq := AnthropicRequest{ @@ -143,7 +143,7 @@ func TestPathTranslationPreservesOtherFields(t *testing.T) { // TestTranslatorSetsPathWithoutOllaPrefix verifies that the translator sets the path correctly WITHOUT /olla prefix // The handler layer is responsible for stripping the /olla prefix, not the translator func TestTranslatorSetsPathWithoutOllaPrefix(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", diff --git a/internal/adapter/translator/anthropic/reasoning_test.go b/internal/adapter/translator/anthropic/reasoning_test.go new file mode 100644 index 00000000..18d31461 --- /dev/null +++ b/internal/adapter/translator/anthropic/reasoning_test.go @@ -0,0 +1,387 @@ +package anthropic + +import ( + "context" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --------------------------------------------------------------------------- +// Streaming tests +// --------------------------------------------------------------------------- + +// TestStreaming_Reasoning_ReasoningField verifies that "reasoning" deltas produce +// the full thinking block sequence before the text block. +func TestStreaming_Reasoning_ReasoningField(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + reasoningChunk("chatcmpl-r1", "deepseek-r1", "Let me think"), + reasoningChunk("chatcmpl-r1", "", " about this"), + textChunk("chatcmpl-r1", "", "The answer is 42"), + finishChunk("chatcmpl-r1", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // thinking block must start before text block + assertThinkingBlockTransitionOrder(t, events) + + // thinking block opened with correct type + starts := findEventsByType(events, "content_block_start") + require.GreaterOrEqual(t, len(starts), 2) + thinkingType, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeThinking, thinkingType, "first block must be thinking") + + textType, _ := getContentBlockType(starts[1]) + assert.Equal(t, contentTypeText, textType, "second block must be text") + + // both blocks must be closed + assertBlocksClosed(t, events) + assertContentBlockCount(t, events, 2) + + // thinking text must appear in thinking_delta events + assert.Contains(t, body, `"type":"thinking_delta"`) + assert.Contains(t, body, `"thinking":"Let me think"`) + assert.Contains(t, body, `"thinking":" about this"`) + + // text must appear in text_delta + assertTextContent(t, body, "The answer is 42") +} + +// TestStreaming_Reasoning_ReasoningContentField verifies that "reasoning_content" +// (vLLM / SGLang / DeepSeek naming) produces identical output to "reasoning". +func TestStreaming_Reasoning_ReasoningContentField(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + reasoningContentChunk("chatcmpl-rc1", "qwq-32b", "Step 1: analyse"), + reasoningContentChunk("chatcmpl-rc1", "", " the problem"), + textChunk("chatcmpl-rc1", "", "Done."), + finishChunk("chatcmpl-rc1", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + assertThinkingBlockTransitionOrder(t, events) + assertBlocksClosed(t, events) + assertContentBlockCount(t, events, 2) + + assert.Contains(t, body, `"type":"thinking_delta"`) + assertThinkingContent(t, body, "Step 1: analyse") + assertThinkingContent(t, body, " the problem") + assertTextContent(t, body, "Done.") +} + +// TestStreaming_Reasoning_OnlyReasoningNoContent verifies that a reasoning-only +// response (no text content) still produces a complete, closed thinking block +// plus the standard message_delta / message_stop tail. +func TestStreaming_Reasoning_OnlyReasoningNoContent(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + reasoningChunk("chatcmpl-ro1", "deepseek-r1", "Internal monologue"), + finishChunk("chatcmpl-ro1", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // exactly one content block: the thinking block + assertContentBlockCount(t, events, 1) + assertBlocksClosed(t, events) + + starts := findEventsByType(events, "content_block_start") + require.Len(t, starts, 1) + bt, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeThinking, bt) + + // full sequence must be present + assertRequiredEvents(t, events) + assert.Contains(t, body, `"thinking":"Internal monologue"`) + assertStopReason(t, events, "end_turn") +} + +// TestStreaming_Reasoning_NoReasoningUnchanged is a regression guard: when there is +// no reasoning field the output must be byte-identical to the pre-reasoning behaviour. +func TestStreaming_Reasoning_NoReasoningUnchanged(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-nr1", "some-model", "Hello world"), + finishChunk("chatcmpl-nr1", "stop"), + doneChunk(), + }) + + recorder := executeTransform(t, tr, stream) + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // no thinking blocks at all + assert.NotContains(t, body, contentTypeThinking) + assert.NotContains(t, body, "thinking_delta") + + assertContentBlockCount(t, events, 1) + starts := findEventsByType(events, "content_block_start") + bt, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeText, bt) + + assertTextContent(t, body, "Hello world") + assertBlocksClosed(t, events) +} + +// TestStreaming_Reasoning_BlockOrderExactSequence pins the full SSE event sequence +// for reasoning + content to ensure the thinking block is fully closed before +// the text block opens (Anthropic spec requirement). +func TestStreaming_Reasoning_BlockOrderExactSequence(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + reasoningChunk("chatcmpl-seq", "deepseek-r1", "thinking..."), + textChunk("chatcmpl-seq", "", "answer"), + finishChunk("chatcmpl-seq", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + events := parseAnthropicEvents(t, recorder.Body.String()) + + assertEventSequence(t, events, []string{ + "message_start", + "content_block_start", // index 0, thinking + "content_block_delta", // thinking_delta + "content_block_stop", // index 0 — thinking closed before text opens + "content_block_start", // index 1, text + "content_block_delta", // text_delta + "content_block_stop", // index 1 + "message_delta", + "message_stop", + }) + + // verify indices + starts := findEventsByType(events, "content_block_start") + idx0, _ := getContentBlockIndex(starts[0]) + assert.Equal(t, 0, idx0) + idx1, _ := getContentBlockIndex(starts[1]) + assert.Equal(t, 1, idx1) +} + +// TestStreaming_Reasoning_MultipleChunksAccumulated verifies that multiple reasoning +// chunks produce separate thinking_delta events (one per chunk) rather than being +// collapsed, preserving the streaming semantics. +func TestStreaming_Reasoning_MultipleChunksAccumulated(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + reasoningChunk("chatcmpl-multi", "deepseek-r1", "chunk1"), + reasoningChunk("chatcmpl-multi", "", "chunk2"), + reasoningChunk("chatcmpl-multi", "", "chunk3"), + textChunk("chatcmpl-multi", "", "result"), + finishChunk("chatcmpl-multi", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // three thinking_delta events (one per chunk), one text_delta + deltas := findEventsByType(events, "content_block_delta") + thinkingDeltas := 0 + textDeltas := 0 + for _, d := range deltas { + delta, ok := d["delta"].(map[string]interface{}) + if !ok { + continue + } + switch delta["type"] { + case "thinking_delta": + thinkingDeltas++ + case "text_delta": + textDeltas++ + } + } + assert.Equal(t, 3, thinkingDeltas, "one thinking_delta per reasoning chunk") + assert.Equal(t, 1, textDeltas, "one text_delta for content") + + // all blocks closed, thinking first + assertBlocksClosed(t, events) + assertThinkingBlockTransitionOrder(t, events) +} + +// --------------------------------------------------------------------------- +// Non-streaming tests +// --------------------------------------------------------------------------- + +// TestResponse_Reasoning_ReasoningField verifies that "reasoning" in a non-streaming +// OpenAI message maps to a leading thinking block in the Anthropic response. +func TestResponse_Reasoning_ReasoningField(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createResponseTestLogger(), createTestConfig()) + + openaiResp := map[string]interface{}{ + "id": "chatcmpl-nr", + "model": "deepseek-r1", + "choices": []interface{}{ + map[string]interface{}{ + "message": map[string]interface{}{ + "role": "assistant", + "reasoning": "First I consider the options.", + "content": "The answer is 7.", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]interface{}{ + "prompt_tokens": float64(10), + "completion_tokens": float64(5), + }, + } + + result, err := tr.TransformResponse(context.Background(), openaiResp, nil) + require.NoError(t, err) + + resp, ok := result.(AnthropicResponse) + require.True(t, ok) + require.Len(t, resp.Content, 2, "thinking block + text block") + + assert.Equal(t, contentTypeThinking, resp.Content[0].Type) + assert.Equal(t, "First I consider the options.", resp.Content[0].Thinking) + assert.Empty(t, resp.Content[0].Text, "thinking block must not set Text") + + assert.Equal(t, contentTypeText, resp.Content[1].Type) + assert.Equal(t, "The answer is 7.", resp.Content[1].Text) +} + +// TestResponse_Reasoning_ReasoningContentField verifies the "reasoning_content" +// field name variant (vLLM / SGLang / DeepSeek) maps to a thinking block. +func TestResponse_Reasoning_ReasoningContentField(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createResponseTestLogger(), createTestConfig()) + + openaiResp := map[string]interface{}{ + "id": "chatcmpl-rcc", + "model": "qwq-32b", + "choices": []interface{}{ + map[string]interface{}{ + "message": map[string]interface{}{ + "role": "assistant", + "reasoning_content": "Step 1: identify the pattern.", + "content": "Pattern identified.", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]interface{}{}, + } + + result, err := tr.TransformResponse(context.Background(), openaiResp, nil) + require.NoError(t, err) + + resp, ok := result.(AnthropicResponse) + require.True(t, ok) + require.Len(t, resp.Content, 2) + + assert.Equal(t, contentTypeThinking, resp.Content[0].Type) + assert.Equal(t, "Step 1: identify the pattern.", resp.Content[0].Thinking) + + assert.Equal(t, contentTypeText, resp.Content[1].Type) + assert.Equal(t, "Pattern identified.", resp.Content[1].Text) +} + +// TestResponse_Reasoning_NoReasoningUnchanged is a regression guard for non-streaming: +// without reasoning fields the content array must not gain a thinking block. +func TestResponse_Reasoning_NoReasoningUnchanged(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createResponseTestLogger(), createTestConfig()) + + openaiResp := map[string]interface{}{ + "id": "chatcmpl-norr", + "model": "llama3", + "choices": []interface{}{ + map[string]interface{}{ + "message": map[string]interface{}{ + "role": "assistant", + "content": "Simple answer.", + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]interface{}{}, + } + + result, err := tr.TransformResponse(context.Background(), openaiResp, nil) + require.NoError(t, err) + + resp, ok := result.(AnthropicResponse) + require.True(t, ok) + require.Len(t, resp.Content, 1, "no thinking block when no reasoning field") + assert.Equal(t, contentTypeText, resp.Content[0].Type) + assert.Equal(t, "Simple answer.", resp.Content[0].Text) +} + +// TestResponse_Reasoning_ReasoningOnlyNoContent verifies that a message with +// reasoning but no content body still yields a valid (single thinking block) response. +func TestResponse_Reasoning_ReasoningOnlyNoContent(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createResponseTestLogger(), createTestConfig()) + + openaiResp := map[string]interface{}{ + "id": "chatcmpl-ron", + "model": "deepseek-r1", + "choices": []interface{}{ + map[string]interface{}{ + "message": map[string]interface{}{ + "role": "assistant", + "reasoning": "Internal thoughts only.", + // no "content" field + }, + "finish_reason": "stop", + }, + }, + "usage": map[string]interface{}{}, + } + + result, err := tr.TransformResponse(context.Background(), openaiResp, nil) + require.NoError(t, err) + + resp, ok := result.(AnthropicResponse) + require.True(t, ok) + + // convertResponseContent falls back to a single empty text block when content is absent; + // with reasoning it should have thinking + the empty text fallback. + thinkingFound := false + for _, block := range resp.Content { + if block.Type == contentTypeThinking { + thinkingFound = true + assert.Equal(t, "Internal thoughts only.", block.Thinking) + } + } + assert.True(t, thinkingFound, "thinking block must be present") +} diff --git a/internal/adapter/translator/anthropic/request.go b/internal/adapter/translator/anthropic/request.go index 71546ad9..38ef9e87 100644 --- a/internal/adapter/translator/anthropic/request.go +++ b/internal/adapter/translator/anthropic/request.go @@ -158,10 +158,12 @@ func (t *Translator) convertSingleMessage(msg AnthropicMessage) ([]map[string]in // user msgs can have text + tool results, assistant msgs have text + tool uses if msg.Role == "user" { userMsg, toolMsgs := t.convertUserMessage(contentBlocks) + // tool messages must immediately follow the assistant tool_calls message; + // OpenAI-compatible backends reject any other role between them. + result = append(result, toolMsgs...) if userMsg != nil { result = append(result, userMsg) } - result = append(result, toolMsgs...) } else if msg.Role == "assistant" { assistantMsg := t.convertAssistantMessage(contentBlocks) if assistantMsg != nil { @@ -203,6 +205,15 @@ func (t *Translator) convertUserMessage(blocks []interface{}) (map[string]interf } } + // OpenAI tool messages have no is_error field; encode the signal in the + // content string so the model can distinguish error results on the way back. + // Skip the prefix if the content already starts with "Error" to avoid + // double-prefixing when re-translating an already-annotated result. + isError, _ := blockMap["is_error"].(bool) + if isError && !strings.HasPrefix(strings.ToLower(content), "error") { + content = "Error: " + content + } + toolResults = append(toolResults, map[string]interface{}{ "role": "tool", "tool_call_id": toolUseID, diff --git a/internal/adapter/translator/anthropic/request_test.go b/internal/adapter/translator/anthropic/request_test.go index 16f13134..2fd62e86 100644 --- a/internal/adapter/translator/anthropic/request_test.go +++ b/internal/adapter/translator/anthropic/request_test.go @@ -21,7 +21,7 @@ func createTestLogger() logger.StyledLogger { } func TestTransformRequest_SimpleMessage(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -60,7 +60,7 @@ func TestTransformRequest_SimpleMessage(t *testing.T) { } func TestTransformRequest_WithSystemPrompt(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -98,7 +98,7 @@ func TestTransformRequest_WithSystemPrompt(t *testing.T) { } func TestTransformRequest_WithTools(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -157,7 +157,7 @@ func TestTransformRequest_WithTools(t *testing.T) { } func TestTransformRequest_MultipleTools(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -212,7 +212,7 @@ func TestTransformRequest_MultipleTools(t *testing.T) { } func TestConvertToolChoice(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) testCases := []struct { name string @@ -276,7 +276,7 @@ func TestConvertToolChoice(t *testing.T) { } func TestConvertToolChoice_EdgeCases(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) t.Run("unknown_string_defaults_to_auto", func(t *testing.T) { result, err := translator.convertToolChoice("unknown") @@ -309,7 +309,7 @@ func TestConvertToolChoice_EdgeCases(t *testing.T) { } func TestConvertMessages_ToolUseAndResult(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -400,7 +400,7 @@ func TestConvertMessages_ToolUseAndResult(t *testing.T) { } func TestTransformRequest_ComplexContent(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -442,7 +442,7 @@ func TestTransformRequest_ComplexContent(t *testing.T) { } func TestTransformRequest_MultipleMessages(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -492,7 +492,7 @@ func TestTransformRequest_MultipleMessages(t *testing.T) { } func TestTransformRequest_EmptyContent(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) t.Run("empty_string_content", func(t *testing.T) { anthropicReq := AnthropicRequest{ @@ -555,7 +555,7 @@ func TestTransformRequest_EmptyContent(t *testing.T) { } func TestTransformRequest_InvalidJSON(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) t.Run("malformed_json", func(t *testing.T) { req := &http.Request{ @@ -578,7 +578,7 @@ func TestTransformRequest_InvalidJSON(t *testing.T) { } func TestTransformRequest_OptionalParameters(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) temp := 0.7 topP := 0.9 @@ -621,7 +621,7 @@ func TestTransformRequest_OptionalParameters(t *testing.T) { } func TestTransformRequest_AssistantWithOnlyToolCalls(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -668,7 +668,7 @@ func TestTransformRequest_AssistantWithOnlyToolCalls(t *testing.T) { } func TestTransformRequest_UserWithOnlyToolResults(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -708,7 +708,7 @@ func TestTransformRequest_UserWithOnlyToolResults(t *testing.T) { } func TestTransformRequest_ToolResultWithStructuredContent(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -759,7 +759,7 @@ func TestTransformRequest_ToolResultWithStructuredContent(t *testing.T) { } func TestTransformRequest_MultipleToolCalls(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -824,7 +824,7 @@ func TestTransformRequest_MultipleToolCalls(t *testing.T) { } func TestConvertToolUse_InvalidData(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) t.Run("missing_id", func(t *testing.T) { block := map[string]interface{}{ @@ -861,7 +861,7 @@ func TestConvertToolUse_InvalidData(t *testing.T) { } func TestTransformRequest_NoMessages(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -882,7 +882,7 @@ func TestTransformRequest_NoMessages(t *testing.T) { } func TestTransformRequest_ToolChoiceObjectForm(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -929,7 +929,7 @@ func TestTransformRequest_ToolChoiceObjectForm(t *testing.T) { } func TestTransformRequest_MixedTextAndToolResults(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -966,16 +966,19 @@ func TestTransformRequest_MixedTextAndToolResults(t *testing.T) { require.True(t, ok) require.Len(t, messages, 2) - assert.Equal(t, "user", messages[0]["role"]) - assert.Equal(t, "Here's the result:", messages[0]["content"]) + // tool messages must come before the user text so they sit immediately after + // the assistant tool_calls message; OpenAI-compatible backends reject any + // other role between them. + assert.Equal(t, "tool", messages[0]["role"]) + assert.Equal(t, "tool_mixed", messages[0]["tool_call_id"]) + assert.Equal(t, "Data from tool", messages[0]["content"]) - assert.Equal(t, "tool", messages[1]["role"]) - assert.Equal(t, "tool_mixed", messages[1]["tool_call_id"]) - assert.Equal(t, "Data from tool", messages[1]["content"]) + assert.Equal(t, "user", messages[1]["role"]) + assert.Equal(t, "Here's the result:", messages[1]["content"]) } func TestConvertSystemPrompt_AllFormats(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) t.Run("string_format", func(t *testing.T) { result := translator.convertSystemPrompt("You are a helpful assistant") @@ -1070,7 +1073,7 @@ func TestConvertSystemPrompt_AllFormats(t *testing.T) { } func TestTransformRequest_SystemPromptWithContentBlocks(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) anthropicReq := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -1115,7 +1118,7 @@ func TestTransformRequest_SystemPromptWithContentBlocks(t *testing.T) { } func TestTransformRequest_StronglyTypedSystemPrompt(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) req := AnthropicRequest{ Model: "claude-3-5-sonnet-20241022", @@ -1150,3 +1153,258 @@ func TestTransformRequest_StronglyTypedSystemPrompt(t *testing.T) { assert.Equal(t, "user", messages[1]["role"]) assert.Equal(t, "What's 2+2?", messages[1]["content"]) } + +// TestConvertUserMessage_ToolResultPrecedesText verifies that when a user turn carries +// both tool_result blocks and a text block (the typical agentic pattern used by Claude +// Code), the resulting OpenAI messages are ordered tool(s) first, user text last. +// +// OpenAI-compatible backends require tool messages to sit immediately after the +// assistant tool_calls message; inserting a user message between them returns 400. +func TestConvertUserMessage_ToolResultPrecedesText(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + anthropicReq := AnthropicRequest{ + Model: "claude-3-5-sonnet-20241022", + MaxTokens: 1024, + Messages: []AnthropicMessage{ + { + Role: "assistant", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_use", + "id": "toolu_agentic", + "name": "bash", + "input": map[string]interface{}{ + "command": "ls", + }, + }, + }, + }, + { + // Claude Code sends [tool_result, text] on virtually every agentic turn. + Role: "user", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "toolu_agentic", + "content": "file1.go file2.go", + }, + map[string]interface{}{ + "type": "text", + "text": "That worked great.", + }, + }, + }, + }, + } + + body, err := json.Marshal(anthropicReq) + require.NoError(t, err) + + req := &http.Request{ + Body: io.NopCloser(bytes.NewReader(body)), + } + + result, err := tr.TransformRequest(context.Background(), req) + require.NoError(t, err) + + messages, ok := result.OpenAIRequest["messages"].([]map[string]interface{}) + require.True(t, ok) + // assistant(tool_calls), tool(result), user(text) + require.Len(t, messages, 3) + + assert.Equal(t, "assistant", messages[0]["role"]) + + assert.Equal(t, "tool", messages[1]["role"], + "tool result must immediately follow the assistant message") + assert.Equal(t, "toolu_agentic", messages[1]["tool_call_id"]) + assert.Equal(t, "file1.go file2.go", messages[1]["content"]) + + assert.Equal(t, "user", messages[2]["role"], + "user text must come after the tool result") + assert.Equal(t, "That worked great.", messages[2]["content"]) +} + +// TestConvertUserMessage_MultipleToolResultsPrecedeText verifies that multiple tool +// results in a single user turn all appear before the user text block. +func TestConvertUserMessage_MultipleToolResultsPrecedeText(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + msgs := []AnthropicMessage{ + { + Role: "user", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "tool_r1", + "content": "result one", + }, + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "tool_r2", + "content": "result two", + }, + map[string]interface{}{ + "type": "text", + "text": "Here is what I got.", + }, + }, + }, + } + + result, err := tr.convertMessages(msgs, nil) + require.NoError(t, err) + require.Len(t, result, 3) + + assert.Equal(t, "tool", result[0]["role"]) + assert.Equal(t, "tool_r1", result[0]["tool_call_id"]) + + assert.Equal(t, "tool", result[1]["role"]) + assert.Equal(t, "tool_r2", result[1]["tool_call_id"]) + + assert.Equal(t, "user", result[2]["role"]) + assert.Equal(t, "Here is what I got.", result[2]["content"]) +} + +// TestConvertUserMessage_TextOnlyUnchanged verifies that a text-only user message +// (no tool_result blocks) is not affected by the reordering. +func TestConvertUserMessage_TextOnlyUnchanged(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + msgs := []AnthropicMessage{ + { + Role: "user", + Content: "Just a plain question.", + }, + } + + result, err := tr.convertMessages(msgs, nil) + require.NoError(t, err) + require.Len(t, result, 1) + + assert.Equal(t, "user", result[0]["role"]) + assert.Equal(t, "Just a plain question.", result[0]["content"]) +} + +// TestConvertUserMessage_IsError_True verifies that a tool_result block with +// is_error:true gets an "Error: " prefix prepended to the content string. +func TestConvertUserMessage_IsError_True(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + msgs := []AnthropicMessage{ + { + Role: "user", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "toolu_err", + "content": "file not found", + "is_error": true, + }, + }, + }, + } + + result, err := tr.convertMessages(msgs, nil) + require.NoError(t, err) + // tool message comes first (before optional user text message) + require.NotEmpty(t, result) + + toolMsg := result[0] + assert.Equal(t, "tool", toolMsg["role"]) + content, _ := toolMsg["content"].(string) + assert.Equal(t, "Error: file not found", content, + "is_error:true must prepend 'Error: ' to the content") +} + +// TestConvertUserMessage_IsError_AlreadyPrefixed verifies that double-prefixing is +// avoided when the content already starts with "Error" (case-insensitive). +func TestConvertUserMessage_IsError_AlreadyPrefixed(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + msgs := []AnthropicMessage{ + { + Role: "user", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "toolu_dup", + "content": "Error: already prefixed message", + "is_error": true, + }, + }, + }, + } + + result, err := tr.convertMessages(msgs, nil) + require.NoError(t, err) + require.NotEmpty(t, result) + + toolMsg := result[0] + content, _ := toolMsg["content"].(string) + assert.Equal(t, "Error: already prefixed message", content, + "content already starting with 'Error' must not be double-prefixed") +} + +// TestConvertUserMessage_IsError_False verifies that is_error:false leaves content unchanged. +func TestConvertUserMessage_IsError_False(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + msgs := []AnthropicMessage{ + { + Role: "user", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "toolu_ok", + "content": "result data", + "is_error": false, + }, + }, + }, + } + + result, err := tr.convertMessages(msgs, nil) + require.NoError(t, err) + require.NotEmpty(t, result) + + toolMsg := result[0] + content, _ := toolMsg["content"].(string) + assert.Equal(t, "result data", content, + "is_error:false must leave content unchanged") +} + +// TestConvertUserMessage_IsError_Absent verifies that omitting is_error leaves content unchanged. +func TestConvertUserMessage_IsError_Absent(t *testing.T) { + t.Parallel() + tr := mustNewTranslator(createTestLogger(), createTestConfig()) + + msgs := []AnthropicMessage{ + { + Role: "user", + Content: []interface{}{ + map[string]interface{}{ + "type": "tool_result", + "tool_use_id": "toolu_absent", + "content": "normal result", + // no is_error field + }, + }, + }, + } + + result, err := tr.convertMessages(msgs, nil) + require.NoError(t, err) + require.NotEmpty(t, result) + + toolMsg := result[0] + content, _ := toolMsg["content"].(string) + assert.Equal(t, "normal result", content, + "absent is_error must leave content unchanged") +} diff --git a/internal/adapter/translator/anthropic/response.go b/internal/adapter/translator/anthropic/response.go index 3e65679e..70890081 100644 --- a/internal/adapter/translator/anthropic/response.go +++ b/internal/adapter/translator/anthropic/response.go @@ -86,6 +86,16 @@ func (t *Translator) TransformResponse(ctx context.Context, openaiResp interface func (t *Translator) convertResponseContent(message map[string]interface{}, finishReason string) ([]ContentBlock, string) { var content []ContentBlock + // Prepend a thinking block when the backend includes chain-of-thought reasoning. + // "reasoning" is used by Ollama/LM Studio/Lemonade; "reasoning_content" by vLLM/SGLang/DeepSeek. + // Local backends do not produce a cryptographic signature, so we omit the signature field. + if reasoning := extractNonStreamingReasoning(message); reasoning != "" { + content = append(content, ContentBlock{ + Type: contentTypeThinking, + Thinking: reasoning, + }) + } + // Handle text content if textContent, ok := message["content"].(string); ok && textContent != "" { content = append(content, ContentBlock{ @@ -170,6 +180,20 @@ func (t *Translator) convertToToolUse(toolCall map[string]interface{}) *ContentB } } +// extractNonStreamingReasoning returns the non-empty reasoning text from an OpenAI message, +// checking both field name variants used across backends: +// - "reasoning" - Ollama, LM Studio, Lemonade +// - "reasoning_content" - vLLM, SGLang, DeepSeek +func extractNonStreamingReasoning(message map[string]interface{}) string { + if v, ok := message["reasoning"].(string); ok && v != "" { + return v + } + if v, ok := message["reasoning_content"].(string); ok && v != "" { + return v + } + return "" +} + // map openai token counts to anthropic names func (t *Translator) convertUsage(usage map[string]interface{}) AnthropicUsage { promptTokens := 0 diff --git a/internal/adapter/translator/anthropic/response_test.go b/internal/adapter/translator/anthropic/response_test.go index 2d3ad94a..4e7cad84 100644 --- a/internal/adapter/translator/anthropic/response_test.go +++ b/internal/adapter/translator/anthropic/response_test.go @@ -16,7 +16,7 @@ func createResponseTestLogger() logger.StyledLogger { } func TestTransformResponse_SimpleText(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "id": "chatcmpl-123", @@ -57,7 +57,7 @@ func TestTransformResponse_SimpleText(t *testing.T) { } func TestTransformResponse_WithToolCalls(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "id": "chatcmpl-123", @@ -109,7 +109,7 @@ func TestTransformResponse_WithToolCalls(t *testing.T) { } func TestTransformResponse_MultipleToolCalls(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "choices": []interface{}{ @@ -158,7 +158,7 @@ func TestTransformResponse_MultipleToolCalls(t *testing.T) { } func TestTransformResponse_EmptyResponse(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("empty_content", func(t *testing.T) { // OpenAI response with empty content string @@ -230,7 +230,7 @@ func TestTransformResponse_EmptyResponse(t *testing.T) { } func TestTransformResponse_MissingUsage(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("usage_not_present", func(t *testing.T) { openaiResp := map[string]interface{}{ @@ -289,7 +289,7 @@ func TestTransformResponse_MissingUsage(t *testing.T) { } func TestTransformResponse_InvalidToolArguments(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("malformed_json", func(t *testing.T) { openaiResp := map[string]interface{}{ @@ -376,7 +376,7 @@ func TestTransformResponse_InvalidToolArguments(t *testing.T) { } func TestTransformResponse_OnlyToolCalls(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "id": "chatcmpl-toolonly", @@ -423,7 +423,7 @@ func TestTransformResponse_OnlyToolCalls(t *testing.T) { } func TestTransformResponse_FinishReasonMapping(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("stop_to_end_turn", func(t *testing.T) { openaiResp := map[string]interface{}{ @@ -562,7 +562,7 @@ func TestTransformResponse_FinishReasonMapping(t *testing.T) { } func TestTransformResponse_MessageIDGeneration(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("generates_anthropic_format_id", func(t *testing.T) { openaiResp := map[string]interface{}{ @@ -665,7 +665,7 @@ func TestTransformResponse_MessageIDGeneration(t *testing.T) { } func TestTransformResponse_ComplexToolArguments(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "id": "chatcmpl-complex", @@ -732,7 +732,7 @@ func TestTransformResponse_ComplexToolArguments(t *testing.T) { } func TestTransformResponse_InvalidFormat(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("missing_choices", func(t *testing.T) { openaiResp := map[string]interface{}{ @@ -792,7 +792,7 @@ func TestTransformResponse_InvalidFormat(t *testing.T) { } func TestTransformResponse_StopSequenceHandling(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) t.Run("with_stop_sequence", func(t *testing.T) { openaiResp := map[string]interface{}{ @@ -821,7 +821,7 @@ func TestTransformResponse_StopSequenceHandling(t *testing.T) { } func TestTransformResponse_ModelFieldPreservation(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) testCases := []struct { name string @@ -870,7 +870,7 @@ func TestTransformResponse_ModelFieldPreservation(t *testing.T) { } func TestTransformResponse_TypeAndRoleFields(t *testing.T) { - translator := NewTranslator(createResponseTestLogger(), createTestConfig()) + translator := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "id": "chatcmpl-fields", @@ -899,7 +899,7 @@ func TestTransformResponse_TypeAndRoleFields(t *testing.T) { } func BenchmarkTransformResponse_Simple(b *testing.B) { - tr := NewTranslator(createResponseTestLogger(), createTestConfig()) + tr := mustNewTranslator(createResponseTestLogger(), createTestConfig()) openaiResp := map[string]interface{}{ "id": "chatcmpl-123", "model": "claude-3-5-sonnet-20241022", "choices": []interface{}{map[string]interface{}{ diff --git a/internal/adapter/translator/anthropic/security_test.go b/internal/adapter/translator/anthropic/security_test.go index a1776935..dee82239 100644 --- a/internal/adapter/translator/anthropic/security_test.go +++ b/internal/adapter/translator/anthropic/security_test.go @@ -14,7 +14,7 @@ import ( ) func TestRequestValidation_RequiredFields(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() t.Run("missing_model", func(t *testing.T) { @@ -59,7 +59,7 @@ func TestRequestValidation_RequiredFields(t *testing.T) { } func TestRequestValidation_ParameterRanges(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() t.Run("negative_max_tokens", func(t *testing.T) { @@ -221,7 +221,7 @@ func TestRequestValidation_ParameterRanges(t *testing.T) { } func TestRequestValidation_ValidParameters(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() t.Run("valid_temperature_boundary", func(t *testing.T) { @@ -292,7 +292,7 @@ func TestRequestValidation_ValidParameters(t *testing.T) { } func TestRequestSizeLimit(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() t.Run("request_exceeding_10MB", func(t *testing.T) { @@ -346,7 +346,7 @@ func TestRequestSizeLimit(t *testing.T) { } func TestUnknownFieldsTolerated(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() t.Run("request_with_unknown_field", func(t *testing.T) { @@ -390,7 +390,7 @@ func TestUnknownFieldsTolerated(t *testing.T) { } func TestSecurityFeaturesCombined(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) ctx := context.Background() t.Run("invalid_params_with_unknown_fields", func(t *testing.T) { diff --git a/internal/adapter/translator/anthropic/streaming.go b/internal/adapter/translator/anthropic/streaming.go index 72cc7a5b..b07d93a9 100644 --- a/internal/adapter/translator/anthropic/streaming.go +++ b/internal/adapter/translator/anthropic/streaming.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -13,6 +14,86 @@ import ( "github.com/thushan/olla/internal/util" ) +// openAIChunk is the typed representation of an incoming OpenAI SSE data line. +// Using a typed struct instead of map[string]interface{} avoids per-chunk map +// allocations and interface boxing - the dominant allocator on the streaming hot path. +type openAIChunk struct { + Usage *openAIUsage `json:"usage"` + Model string `json:"model"` + Choices []openAIChoice `json:"choices"` +} + +type openAIChoice struct { + FinishReason *string `json:"finish_reason"` + Delta openAIDelta `json:"delta"` + Index int `json:"index"` +} + +type openAIDelta struct { + Role string `json:"role"` + Content *string `json:"content"` + Reasoning string `json:"reasoning"` + ReasoningContent string `json:"reasoning_content"` + ToolCalls []openAIToolCall `json:"tool_calls"` +} + +type openAIToolCall struct { + Index *int `json:"index"` + ID string `json:"id"` + Type string `json:"type"` + Function openAIFunction `json:"function"` +} + +type openAIFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +// openAIUsage holds token counts from the final SSE chunk. +// Fields are int rather than float64 because json.Unmarshal into a typed struct +// uses the declared field type directly, avoiding the float64-then-int cast +// that the map[string]interface{} path required. +type openAIUsage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` +} + +// sseDeltaEvent is the hot-path content_block_delta envelope. +// Field order is chosen for struct packing (interface{} + string + int fits without +// padding on amd64). JSON tag order (delta, index, type) differs from field +// declaration order; clients parse JSON by key, not position, so this is safe. +type sseDeltaEvent struct { + Delta interface{} `json:"delta"` // sseTextDelta | sseThinkingDelta | sseInputJSONDelta + Type string `json:"type"` + Index int `json:"index"` +} + +// sseTextDelta is the delta payload for text_delta events. +// Field order matches alphabetical JSON tag order (text < type). +type sseTextDelta struct { + Text string `json:"text"` + Type string `json:"type"` +} + +// sseThinkingDelta is the delta payload for thinking_delta events. +// Field order matches alphabetical JSON tag order (thinking < type). +type sseThinkingDelta struct { + Thinking string `json:"thinking"` + Type string `json:"type"` +} + +// sseInputJSONDelta is the delta payload for input_json_delta events. +// Field order matches alphabetical JSON tag order (partial_json < type). +type sseInputJSONDelta struct { + PartialJSON string `json:"partial_json"` + Type string `json:"type"` +} + +// errInterleavedToolArguments signals that args arrived for a tool block that was +// already stopped and closed on the wire. Corrupt tool input is worse than a failed +// stream; the client can retry a clean failure, but cannot recover silent data loss. +var errInterleavedToolArguments = errors.New("tool arguments received for already-stopped block") + // tracks state while streaming - buffers partial data, blocks in progress type StreamingState struct { currentBlock *ContentBlock @@ -26,6 +107,9 @@ type StreamingState struct { inputTokens int outputTokens int messageStartSent bool + // reasoningOpen is true while a thinking block has been started but not yet stopped. + // The first non-reasoning delta closes the thinking block before opening its own. + reasoningOpen bool } // convert openai sse stream to anthropic format @@ -44,10 +128,21 @@ func (t *Translator) TransformStreamingResponse(ctx context.Context, openaiStrea toolIndexToBlock: make(map[int]int), } + // Seed input_tokens from the pre-computed estimate injected by the handler. + // vLLM and lmdeploy both emit real input_tokens in message_start rather than 0; + // this brings the translation path in line with that behaviour. The value is + // overwritten by actual upstream usage when (and if) the backend sends it. + if estimate, ok := ctx.Value(constants.ContextInputTokensKey).(int); ok && estimate > 0 { + state.inputTokens = estimate + } + // sync streaming for now (async needs more work for agent workflows) streamErr := t.transformStreamingSync(ctx, openaiStream, w, rc, state) if streamErr != nil { + // Return immediately for all stream errors. For errInterleavedToolArguments the + // SSE error event is already on the wire; emitting message_delta/message_stop after + // an error event is invalid per the Anthropic spec. return streamErr } @@ -87,6 +182,10 @@ func (t *Translator) transformStreamingSync(ctx context.Context, openaiStream io line := scanner.Text() if err := t.processStreamLine(line, state, w, rc); err != nil { + // Interleaved tool args corrupt the client's tool input; abort the stream. + if errors.Is(err, errInterleavedToolArguments) { + return err + } t.logger.Error("Error processing stream line", "error", err) continue // keep going, don't fail entire stream on one bad line } @@ -110,7 +209,7 @@ func (t *Translator) processStreamLine(line string, state *StreamingState, w htt return nil } - var chunk map[string]interface{} + var chunk openAIChunk if err := json.Unmarshal([]byte(data), &chunk); err != nil { // log bad chunks but keep going, partial responses better than nothing t.logger.Warn("Malformed chunk encountered, skipping", "error", err, @@ -119,53 +218,74 @@ func (t *Translator) processStreamLine(line string, state *StreamingState, w htt } // grab model name for message_start event - if state.model == "" { - if model, ok := chunk["model"].(string); ok { - state.model = model - } + if state.model == "" && chunk.Model != "" { + state.model = chunk.Model } - choices, ok := chunk["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil + // Capture usage before the empty-choices guard. vLLM (and others that honour + // include_usage=true) emit the final usage as a standalone chunk with + // choices:[] and a populated usage object. Processing choices first and + // returning early would silently discard those real token counts, leaving + // message_delta with the pre-stream estimate or output_tokens:0. + if u := chunk.Usage; u != nil { + state.inputTokens = u.PromptTokens + state.outputTokens = u.CompletionTokens } - choice, ok := choices[0].(map[string]interface{}) - if !ok { + if len(chunk.Choices) == 0 { return nil } - // capture finish_reason for later stop_reason mapping - if finishReason, finishOk := choice["finish_reason"].(string); finishOk && finishReason != "" { - state.lastFinishReason = finishReason - } + choice := chunk.Choices[0] - // grab usage stats if present (usually in final chunk) - if usage, usageOk := chunk["usage"].(map[string]interface{}); usageOk { - if promptTokens, promptOk := usage["prompt_tokens"].(float64); promptOk { - state.inputTokens = int(promptTokens) - } - if completionTokens, completionsOk := usage["completion_tokens"].(float64); completionsOk { - state.outputTokens = int(completionTokens) - } + // capture finish_reason for later stop_reason mapping; nil or empty means absent + if fr := choice.FinishReason; fr != nil && *fr != "" { + state.lastFinishReason = *fr } - delta, ok := choice["delta"].(map[string]interface{}) - if !ok { - return nil + delta := choice.Delta + + // reasoning / reasoning_content are the per-backend field names for chain-of-thought. + // Ollama, LM Studio, and Lemonade use "reasoning"; vLLM, SGLang, and DeepSeek use + // "reasoning_content". Treat them as equivalent and prefer whichever is non-empty. + // + // Some OpenAI-compatible adapters coalesce reasoning + content (or + tool_calls) into + // a single chunk. We must not return early after handling reasoning: fall through so + // content and tool_calls in the same chunk are also processed. handleContentDelta and + // initializeToolBlock both call closeCurrentBlock before opening a new one, so the + // thinking block is correctly stopped before the text or tool block opens. + reasoning := extractReasoningFieldTyped(delta) + if reasoning != "" { + if err := t.handleReasoningDelta(reasoning, state, w, rc); err != nil { + return err + } } - if content, ok := delta["content"].(string); ok && content != "" { - return t.handleContentDelta(content, state, w, rc) + // Content is a pointer so we can distinguish absent/null (nil) from empty string. + if delta.Content != nil && *delta.Content != "" { + if err := t.handleContentDelta(*delta.Content, state, w, rc); err != nil { + return err + } } - if toolCalls, ok := delta["tool_calls"].([]interface{}); ok { - return t.handleToolCallsDelta(toolCalls, state, w, rc) + if len(delta.ToolCalls) > 0 { + return t.handleToolCallsDelta(delta.ToolCalls, state, w, rc) } return nil } +// extractReasoningFieldTyped returns the non-empty reasoning text from a typed delta, +// checking both field name variants used across backends: +// - "reasoning" - Ollama, LM Studio, Lemonade +// - "reasoning_content" - vLLM, SGLang, DeepSeek +func extractReasoningFieldTyped(delta openAIDelta) string { + if delta.Reasoning != "" { + return delta.Reasoning + } + return delta.ReasoningContent +} + // send message_start if we haven't already, needs to be first event func (t *Translator) ensureMessageStartSent(state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { if !state.messageStartSent { @@ -188,17 +308,10 @@ func (t *Translator) handleContentDelta(content string, state *StreamingState, w // start new text block if needed (anthropic wants block_start before deltas) if state.currentBlock == nil || state.currentBlock.Type != contentTypeText { - // close previous block if different type - if state.currentBlock != nil && state.currentBlock.Type != contentTypeText { - if err := t.writeEvent(w, "content_block_stop", map[string]interface{}{ - "type": "content_block_stop", - "index": state.currentIndex, - }); err != nil { - return err - } - if err := rc.Flush(); err != nil { - return fmt.Errorf("flush failed: %w", err) - } + // close whatever block is open (tool_use or a prior text block) before + // opening a new text block; every start needs a matching stop. + if err := t.closeCurrentBlock(state, w, rc); err != nil { + return err } state.currentBlock = &ContentBlock{ @@ -223,14 +336,11 @@ func (t *Translator) handleContentDelta(content string, state *StreamingState, w } } - // send delta event for each chunk - if err := t.writeEvent(w, "content_block_delta", map[string]interface{}{ - "type": "content_block_delta", - "index": state.currentIndex, - "delta": map[string]interface{}{ - "type": "text_delta", - "text": content, - }, + // send delta event for each chunk - typed struct avoids a per-chunk map allocation + if err := t.writeEvent(w, "content_block_delta", sseDeltaEvent{ + Delta: sseTextDelta{Text: content, Type: "text_delta"}, + Index: state.currentIndex, + Type: "content_block_delta", }); err != nil { return err } @@ -245,53 +355,36 @@ func (t *Translator) handleContentDelta(content string, state *StreamingState, w return nil } -// toolCallData holds extracted and validated tool call information -type toolCallData struct { - id string - name string - arguments string - toolIndex int -} - -// extractToolCallData validates and extracts data from a tool call delta -func extractToolCallData(tc interface{}) (*toolCallData, bool) { - toolCall, ok := tc.(map[string]interface{}) - if !ok { - return nil, false - } - - index, _ := toolCall["index"].(float64) - toolIndex := int(index) - - function, ok := toolCall["function"].(map[string]interface{}) - if !ok { - return nil, false - } - - data := &toolCallData{ - toolIndex: toolIndex, +// handleReasoningDelta streams chain-of-thought text as an Anthropic thinking block. +// Reasoning always precedes content in model output, so we open a thinking block on the +// first reasoning chunk and leave it open until the first non-reasoning delta arrives, +// at which point closeCurrentBlock handles the stop event before the next block opens. +func (t *Translator) handleReasoningDelta(reasoning string, state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { + if err := t.ensureMessageStartSent(state, w, rc); err != nil { + return err } - // extract optional fields - if id, ok := toolCall["id"].(string); ok { - data.id = id - } - if name, ok := function["name"].(string); ok { - data.name = name - } - if args, ok := function["arguments"].(string); ok { - data.arguments = args - } + // Open the thinking block on the first reasoning chunk. + if !state.reasoningOpen { + // Close anything that was already open (shouldn't happen in practice, but be safe). + if err := t.closeCurrentBlock(state, w, rc); err != nil { + return err + } - return data, true -} + state.currentBlock = &ContentBlock{ + Type: contentTypeThinking, + } + state.currentIndex = len(state.contentBlocks) + state.contentBlocks = append(state.contentBlocks, *state.currentBlock) + state.reasoningOpen = true -// closeCurrentBlockIfNeeded closes the current block if it exists and matches the given type -func (t *Translator) closeCurrentBlockIfNeeded(state *StreamingState, blockType string, w http.ResponseWriter, rc *http.ResponseController) error { - if state.currentBlock != nil && state.currentBlock.Type == blockType { - if err := t.writeEvent(w, "content_block_stop", map[string]interface{}{ - "type": "content_block_stop", + if err := t.writeEvent(w, "content_block_start", map[string]interface{}{ + "type": "content_block_start", "index": state.currentIndex, + "content_block": map[string]interface{}{ + "type": contentTypeThinking, + "thinking": "", + }, }); err != nil { return err } @@ -299,13 +392,60 @@ func (t *Translator) closeCurrentBlockIfNeeded(state *StreamingState, blockType return fmt.Errorf("flush failed: %w", err) } } + + // typed struct avoids a per-chunk map allocation on the thinking_delta hot path + if err := t.writeEvent(w, "content_block_delta", sseDeltaEvent{ + Delta: sseThinkingDelta{Thinking: reasoning, Type: "thinking_delta"}, + Index: state.currentIndex, + Type: "content_block_delta", + }); err != nil { + return err + } + if err := rc.Flush(); err != nil { + return fmt.Errorf("flush failed: %w", err) + } + + // Accumulate for inspector/logging. + state.currentBlock.Thinking += reasoning + state.contentBlocks[state.currentIndex] = *state.currentBlock + return nil } -// initializeToolBlock creates and sends a new tool_use block start event +// closeCurrentBlock closes the currently open block regardless of type. +// Every content_block_start must be paired with a content_block_stop; SDK +// accumulators (including Claude Code's own stream reader) finalise tool +// input on the stop event, so omitting it silently breaks multi-tool calls. +func (t *Translator) closeCurrentBlock(state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { + if state.currentBlock == nil { + return nil + } + if err := t.writeEvent(w, "content_block_stop", map[string]interface{}{ + "type": "content_block_stop", + "index": state.currentIndex, + }); err != nil { + return err + } + if err := rc.Flush(); err != nil { + return fmt.Errorf("flush failed: %w", err) + } + // Clear thinking state when closing a thinking block so subsequent reasoning + // chunks (if any) would open a new block rather than append to a closed one. + if state.currentBlock.Type == contentTypeThinking { + state.reasoningOpen = false + } + state.currentBlock = nil + return nil +} + +// initializeToolBlock creates and sends a new tool_use block start event. +// If args arrived before the id+name chunk (pre-init buffering), flushes the +// buffer as a single input_json_delta immediately after content_block_start so +// the pre-init case is lossless. func (t *Translator) initializeToolBlock(id, name string, toolIndex int, state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { - // close current text block before starting tool block, anthropic requires this - if err := t.closeCurrentBlockIfNeeded(state, contentTypeText, w, rc); err != nil { + // close whatever block is currently open before starting a new one; both text + // and prior tool blocks must be stopped before the next block_start is emitted. + if err := t.closeCurrentBlock(state, w, rc); err != nil { return err } @@ -332,20 +472,68 @@ func (t *Translator) initializeToolBlock(id, name string, toolIndex int, state * return err } - return rc.Flush() + if err := rc.Flush(); err != nil { + return fmt.Errorf("flush failed: %w", err) + } + + // Flush any args that arrived before the id+name chunk opened this block. + // Without this the pre-init buffer is silently discarded on the wire. + if buf, exists := state.toolCallBuffers[toolIndex]; exists && buf.Len() > 0 { + buffered := buf.String() + if err := t.writeEvent(w, "content_block_delta", sseDeltaEvent{ + Delta: sseInputJSONDelta{PartialJSON: buffered, Type: "input_json_delta"}, + Index: state.currentIndex, + Type: "content_block_delta", + }); err != nil { + return err + } + if err := rc.Flush(); err != nil { + return fmt.Errorf("flush failed: %w", err) + } + } + + return nil } -// sendToolArgumentsDelta buffers and sends a partial_json delta for tool arguments +// sendToolArgumentsDelta buffers and emits a partial_json delta for the given tool index. +// +// Three cases: +// - Block not yet initialised: buffer only; initializeToolBlock will flush on init. +// - Block is the current open block: emit input_json_delta normally. +// - Block was already stopped: corrupt wire state - emit an SSE error event and +// return errInterleavedToolArguments so the stream is aborted cleanly. func (t *Translator) sendToolArgumentsDelta(args string, toolIndex int, state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { + // Always buffer regardless of block state so finalisation stays correct. state.toolCallBuffers[toolIndex].WriteString(args) - if err := t.writeEvent(w, "content_block_delta", map[string]interface{}{ - "type": "content_block_delta", - "index": state.currentIndex, - "delta": map[string]interface{}{ - "type": "input_json_delta", - "partial_json": args, - }, + blockIndex, mapped := state.toolIndexToBlock[toolIndex] + if !mapped { + // Args arrived before id+name chunk; initializeToolBlock flushes the buffer. + t.logger.Debug("Tool args received before block initialised, buffering only", + "tool_index", toolIndex) + return nil + } + + // Interleaved or late delivery: args for a block that is already stopped. + // Corrupt tool input is worse than a failed stream - abort with a spec-valid error event. + if state.currentBlock == nil || blockIndex != state.currentIndex { + msg := fmt.Sprintf("tool arguments received for already-stopped block (tool_index=%d, block_index=%d)", toolIndex, blockIndex) + _ = t.writeEvent(w, "error", map[string]interface{}{ + "type": "error", + "error": map[string]interface{}{ + "type": "api_error", + "message": msg, + }, + }) + _ = rc.Flush() + return errInterleavedToolArguments + } + + // typed struct avoids a per-chunk map allocation on the input_json_delta hot path + if err := t.writeEvent(w, "content_block_delta", sseDeltaEvent{ + Delta: sseInputJSONDelta{PartialJSON: args, Type: "input_json_delta"}, + Index: blockIndex, + Type: "content_block_delta", }); err != nil { return err } @@ -354,32 +542,41 @@ func (t *Translator) sendToolArgumentsDelta(args string, toolIndex int, state *S } // process tool call deltas, buffers partial json args -func (t *Translator) handleToolCallsDelta(toolCalls []interface{}, state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { +func (t *Translator) handleToolCallsDelta(toolCalls []openAIToolCall, state *StreamingState, w http.ResponseWriter, rc *http.ResponseController) error { if err := t.ensureMessageStartSent(state, w, rc); err != nil { return err } for _, tc := range toolCalls { - data, ok := extractToolCallData(tc) - if !ok { - continue + // absent index field defaults to 0 - the first and usually only tool call + toolIndex := 0 + if tc.Index != nil { + toolIndex = *tc.Index } // initialise buffer if first time seeing this tool index - if _, exists := state.toolCallBuffers[data.toolIndex]; !exists { - state.toolCallBuffers[data.toolIndex] = &strings.Builder{} + if _, exists := state.toolCallBuffers[toolIndex]; !exists { + state.toolCallBuffers[toolIndex] = &strings.Builder{} } - // start block when we get id + name - if data.id != "" && data.name != "" { - if err := t.initializeToolBlock(data.id, data.name, data.toolIndex, state, w, rc); err != nil { - return err + // Start a new block when we receive id + name for the first time for this tool index. + // Some backends repeat id + name on every continuation chunk alongside the args; + // re-initialising would close the live block and open a duplicate, corrupting the + // tool input. The toolIndexToBlock mapping is the authoritative "already started" signal. + if tc.ID != "" && tc.Function.Name != "" { + if _, alreadyInit := state.toolIndexToBlock[toolIndex]; !alreadyInit { + if err := t.initializeToolBlock(tc.ID, tc.Function.Name, toolIndex, state, w, rc); err != nil { + return err + } + } else { + t.logger.Debug("Ignoring repeated id+name for already-initialised tool block", + "tool_index", toolIndex, "tool_id", tc.ID, "tool_name", tc.Function.Name) } } // buffer args chunks and send as partial_json - if data.arguments != "" { - if err := t.sendToolArgumentsDelta(data.arguments, data.toolIndex, state, w, rc); err != nil { + if tc.Function.Arguments != "" { + if err := t.sendToolArgumentsDelta(tc.Function.Arguments, toolIndex, state, w, rc); err != nil { return err } } @@ -390,17 +587,10 @@ func (t *Translator) handleToolCallsDelta(toolCalls []interface{}, state *Stream // send final events, parse tool buffers, determine stop_reason func (t *Translator) finalizeStream(state *StreamingState, w http.ResponseWriter, rc *http.ResponseController, original *http.Request) error { - // close current block if still open - if state.currentBlock != nil { - if err := t.writeEvent(w, "content_block_stop", map[string]interface{}{ - "type": "content_block_stop", - "index": state.currentIndex, - }); err != nil { - return err - } - if err := rc.Flush(); err != nil { - return fmt.Errorf("flush failed: %w", err) - } + // close the last open block; with the single-active-block invariant restored + // by closeCurrentBlock, exactly one stop is emitted here at most. + if err := t.closeCurrentBlock(state, w, rc); err != nil { + return err } // parse buffered json args into objects using the tool index mapping @@ -433,6 +623,27 @@ func (t *Translator) finalizeStream(state *StreamingState, w http.ResponseWriter // map finish_reason to stop_reason (same logic as non-streaming) stopReason := mapFinishReasonToStopReason(state.lastFinishReason) + // When the backend stream carries no usage (common for OpenAI-compatible backends + // that don't set include_usage), output_tokens stays 0 from initialisation. + // Synthesise an estimate from accumulated content so Anthropic clients (including + // Claude Code) receive a plausible non-zero value rather than 0. + // The chars/4 heuristic matches the count_tokens estimator - acceptable precision + // for a fallback. Real backend usage always takes precedence; this branch only fires + // when state.outputTokens is still 0 after the full stream has been consumed. + if state.outputTokens == 0 { + var charCount int + for _, block := range state.contentBlocks { + charCount += len(block.Text) + len(block.Thinking) + } + if charCount > 0 { + estimated := charCount / 4 + if estimated < 1 { + estimated = 1 + } + state.outputTokens = estimated + } + } + // send delta with stop_reason + usage if err := t.writeEvent(w, "message_delta", map[string]interface{}{ "type": "message_delta", @@ -441,8 +652,10 @@ func (t *Translator) finalizeStream(state *StreamingState, w http.ResponseWriter "stop_sequence": nil, }, "usage": map[string]interface{}{ - "input_tokens": state.inputTokens, - "output_tokens": state.outputTokens, + "input_tokens": state.inputTokens, + "output_tokens": state.outputTokens, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, }, }); err != nil { return err @@ -513,29 +726,64 @@ func (t *Translator) createMessageStart(state *StreamingState) map[string]interf return map[string]interface{}{ "type": "message_start", "message": map[string]interface{}{ - "id": state.messageID, - "type": "message", - "role": "assistant", - "model": state.model, - "content": []interface{}{}, + "id": state.messageID, + "type": "message", + "role": "assistant", + "model": state.model, + "content": []interface{}{}, + "stop_reason": nil, + "stop_sequence": nil, "usage": map[string]interface{}{ - "input_tokens": state.inputTokens, - "output_tokens": 0, + "input_tokens": state.inputTokens, + "output_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, }, }, } } -// write sse event: event: \ndata: \n\n +// writeEvent serialises data as a Server-Sent Events frame and writes it to w. +// +// Uses the translator's buffer pool to avoid a heap allocation per event. +// The SSE frame format is: "event: \ndata: \n\n". +// The pooled buffer is reset and returned before this function returns, so +// there is no aliasing risk - w receives a copy via its own internal write path. func (t *Translator) writeEvent(w http.ResponseWriter, event string, data interface{}) error { - dataJSON, err := json.Marshal(data) + buf, err := t.bufferPool.Get() if err != nil { + return fmt.Errorf("writeEvent: buffer pool exhausted: %w", err) + } + // Encode JSON directly into the pooled buffer, avoiding a separate []byte allocation. + enc := json.NewEncoder(buf) + if err := enc.Encode(data); err != nil { + t.bufferPool.Put(buf) return fmt.Errorf("failed to marshal event data: %w", err) } + // json.Encoder.Encode appends a trailing newline; trim it so our framing is exact. + dataJSON := buf.Bytes() + if len(dataJSON) > 0 && dataJSON[len(dataJSON)-1] == '\n' { + dataJSON = dataJSON[:len(dataJSON)-1] + } - if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, dataJSON); err != nil { - return fmt.Errorf("failed to write event: %w", err) + // Write the SSE frame in three direct writes - cheaper than fmt.Fprintf's + // format string parsing and avoids a string concatenation allocation. + var writeErr error + if _, writeErr = io.WriteString(w, "event: "); writeErr == nil { + if _, writeErr = io.WriteString(w, event); writeErr == nil { + if _, writeErr = io.WriteString(w, "\ndata: "); writeErr == nil { + if _, writeErr = w.Write(dataJSON); writeErr == nil { + _, writeErr = io.WriteString(w, "\n\n") + } + } + } } + buf.Reset() + t.bufferPool.Put(buf) + + if writeErr != nil { + return fmt.Errorf("failed to write event: %w", writeErr) + } return nil } diff --git a/internal/adapter/translator/anthropic/streaming_bench_test.go b/internal/adapter/translator/anthropic/streaming_bench_test.go new file mode 100644 index 00000000..12e77853 --- /dev/null +++ b/internal/adapter/translator/anthropic/streaming_bench_test.go @@ -0,0 +1,238 @@ +package anthropic + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" +) + +// buildTextStream constructs a synthetic OpenAI SSE stream with n text delta chunks +// followed by a finish chunk and the done marker. Used to exercise the content_block_delta +// hot path without the scanner/parse overhead dominating the benchmark. +func buildTextStream(n int, chunkText string) io.Reader { + var sb strings.Builder + // First chunk carries model name. + sb.WriteString(textChunk("chatcmpl-bench", "claude-3-5-sonnet-20241022", chunkText)) + for i := 1; i < n; i++ { + sb.WriteString(textChunk("chatcmpl-bench", "", chunkText)) + } + sb.WriteString(finishChunk("chatcmpl-bench", "stop")) + sb.WriteString(doneChunk()) + return strings.NewReader(sb.String()) +} + +// buildThinkingStream constructs a synthetic stream with n reasoning delta chunks +// followed by a single text chunk, finish, and done. +func buildThinkingStream(n int, chunkText string) io.Reader { + var sb strings.Builder + sb.WriteString(reasoningChunk("chatcmpl-bench-r", "deepseek-r1", chunkText)) + for i := 1; i < n; i++ { + sb.WriteString(reasoningChunk("chatcmpl-bench-r", "", chunkText)) + } + sb.WriteString(textChunk("chatcmpl-bench-r", "", "result")) + sb.WriteString(finishChunk("chatcmpl-bench-r", "stop")) + sb.WriteString(doneChunk()) + return strings.NewReader(sb.String()) +} + +// buildToolStream constructs a stream with a single tool call whose arguments arrive +// in n small chunks, exercising the input_json_delta hot path. +func buildToolStream(n int) io.Reader { + var sb strings.Builder + sb.WriteString(toolStartChunk("chatcmpl-bench-t", 0, "call_bench", "bench_tool")) + for i := range n { + sb.WriteString(toolArgsChunk("chatcmpl-bench-t", 0, fmt.Sprintf(`{"n":%d}`, i))) + } + sb.WriteString(finishChunk("chatcmpl-bench-t", "tool_calls")) + sb.WriteString(doneChunk()) + return strings.NewReader(sb.String()) +} + +// BenchmarkStreaming_ContentDeltas measures allocations on the text_delta hot path. +// Run with: go test -bench=BenchmarkStreaming_ContentDeltas -benchmem -count=5 +func BenchmarkStreaming_ContentDeltas(b *testing.B) { + const chunks = 50 // representative of a medium-length response + + tr := newTestTranslator() + ctx := context.Background() + + b.ReportAllocs() + b.ResetTimer() + + for range b.N { + b.StopTimer() + stream := buildTextStream(chunks, "hello ") + rec := httptest.NewRecorder() + b.StartTimer() + + if err := tr.TransformStreamingResponse(ctx, stream, rec, nil); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkStreaming_ThinkingDeltas measures allocations on the thinking_delta hot path. +// Run with: go test -bench=BenchmarkStreaming_ThinkingDeltas -benchmem -count=5 +func BenchmarkStreaming_ThinkingDeltas(b *testing.B) { + const chunks = 50 + + tr := newTestTranslator() + ctx := context.Background() + + b.ReportAllocs() + b.ResetTimer() + + for range b.N { + b.StopTimer() + stream := buildThinkingStream(chunks, "thinking step ") + rec := httptest.NewRecorder() + b.StartTimer() + + if err := tr.TransformStreamingResponse(ctx, stream, rec, nil); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkStreaming_ToolArgDeltas measures allocations on the input_json_delta hot path. +// Run with: go test -bench=BenchmarkStreaming_ToolArgDeltas -benchmem -count=5 +func BenchmarkStreaming_ToolArgDeltas(b *testing.B) { + const chunks = 50 + + tr := newTestTranslator() + ctx := context.Background() + + b.ReportAllocs() + b.ResetTimer() + + for range b.N { + b.StopTimer() + stream := buildToolStream(chunks) + rec := httptest.NewRecorder() + b.StartTimer() + + if err := tr.TransformStreamingResponse(ctx, stream, rec, nil); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkWriteEvent_TextDelta isolates the writeEvent + JSON encode path for a single +// text_delta event. This directly measures the allocation savings from the pooled buffer +// and typed struct, independent of the incoming OpenAI chunk parse overhead. +// Run with: go test -bench=BenchmarkWriteEvent -benchmem -count=5 +func BenchmarkWriteEvent_TextDelta(b *testing.B) { + tr := newTestTranslator() + rec := httptest.NewRecorder() + evt := sseDeltaEvent{ + Delta: sseTextDelta{Text: "hello world", Type: "text_delta"}, + Index: 0, + Type: "content_block_delta", + } + + b.ReportAllocs() + b.ResetTimer() + + for range b.N { + rec.Body.Reset() + if err := tr.writeEvent(rec, "content_block_delta", evt); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkWriteEvent_TextDelta_MapBaseline is a reference benchmark that uses the old +// map[string]interface{} approach so the before/after is visible in a single run. +// It is not wired into any production path - delete after measurements are taken. +func BenchmarkWriteEvent_TextDelta_MapBaseline(b *testing.B) { + tr := newTestTranslator() + rec := httptest.NewRecorder() + + b.ReportAllocs() + b.ResetTimer() + + for range b.N { + rec.Body.Reset() + // Replicate exact old writeEvent logic: json.Marshal on a nested map + fmt.Fprintf. + data := map[string]interface{}{ + "type": "content_block_delta", + "index": 0, + "delta": map[string]interface{}{ + "type": "text_delta", + "text": "hello world", + }, + } + if err := tr.writeEvent(rec, "content_block_delta", data); err != nil { + b.Fatal(err) + } + } +} + +// TestWriteEvent_ConcurrentPoolSafety exercises writeEvent from multiple goroutines +// simultaneously to confirm the pooled buffer is not aliased across callers. +// Run with: go test -run TestWriteEvent_ConcurrentPoolSafety -race +func TestWriteEvent_ConcurrentPoolSafety(t *testing.T) { + t.Parallel() + + tr := newTestTranslator() + ctx := context.Background() + + const goroutines = 20 + const chunksEach = 30 + + var wg sync.WaitGroup + wg.Add(goroutines) + + for i := range goroutines { + go func(id int) { + defer wg.Done() + stream := buildTextStream(chunksEach, fmt.Sprintf("goroutine %d ", id)) + rec := httptest.NewRecorder() + if err := tr.TransformStreamingResponse(ctx, stream, rec, nil); err != nil { + t.Errorf("goroutine %d: unexpected error: %v", id, err) + } + }(i) + } + + wg.Wait() +} + +// BenchmarkProcessStreamLine_TextChunk isolates the per-line parse + dispatch cost +// for a single text delta chunk, after message_start has already been sent. +// This directly measures the typed-struct decode savings on the hot path, independent +// of scanner, HTTP, and SSE write overhead. +// Run with: go test -bench=BenchmarkProcessStreamLine_TextChunk -benchmem -count=5 +func BenchmarkProcessStreamLine_TextChunk(b *testing.B) { + tr := newTestTranslator() + + line := `data: {"id":"chatcmpl-b","model":"claude-3-5-sonnet-20241022","choices":[{"delta":{"content":"hello"},"index":0}]}` + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + // Recreate recorder and state each iteration so alloc counts reflect a + // single-call steady state and mutable fields don't bleed between runs. + rec := httptest.NewRecorder() + rc := http.NewResponseController(rec) + state := &StreamingState{ + messageID: "msg-bench", + contentBlocks: make([]ContentBlock, 0, 4), + toolCallBuffers: make(map[int]*strings.Builder), + toolIndexToBlock: make(map[int]int), + messageStartSent: true, + } + // Prime a text block so the hot path skips block_start overhead. + state.currentBlock = &ContentBlock{Type: contentTypeText, Text: ""} + state.currentIndex = 0 + state.contentBlocks = append(state.contentBlocks, *state.currentBlock) + + if err := tr.processStreamLine(line, state, rec, rc); err != nil { + b.Fatal(err) + } + } +} diff --git a/internal/adapter/translator/anthropic/streaming_parse_test.go b/internal/adapter/translator/anthropic/streaming_parse_test.go new file mode 100644 index 00000000..38247f41 --- /dev/null +++ b/internal/adapter/translator/anthropic/streaming_parse_test.go @@ -0,0 +1,474 @@ +package anthropic + +import ( + "context" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestProcessStreamLine_NullContentDelta verifies that a delta with an explicit +// "content":null value does not emit any content_block_delta event. JSON null +// unmarshals to nil for a *string, so the nil-check on the typed path must hold. +func TestProcessStreamLine_NullContentDelta(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // content is explicitly null — not absent, not empty string + stream := mockOpenAIStream([]string{ + `data: {"id":"chatcmpl-null","model":"claude-3-5-sonnet-20241022","choices":[{"delta":{"content":null},"index":0}]}` + "\n\n", + finishChunk("chatcmpl-null", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + // null content must not produce any content_block_delta + assertEventCount(t, events, "content_block_delta", 0) + // but the stream must still complete + assertHasEventType(t, events, "message_start") + assertHasEventType(t, events, "message_stop") +} + +// TestProcessStreamLine_RoleOnlyDelta verifies that a delta carrying only a +// "role":"assistant" field emits no content blocks at all. This is the first +// chunk pattern emitted by some backends to announce the role before content +// begins streaming. +func TestProcessStreamLine_RoleOnlyDelta(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + `data: {"id":"chatcmpl-role","model":"claude-3-5-sonnet-20241022","choices":[{"delta":{"role":"assistant"},"index":0}]}` + "\n\n", + finishChunk("chatcmpl-role", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + // role-only delta must produce zero content blocks and zero deltas + assertEventCount(t, events, "content_block_start", 0) + assertEventCount(t, events, "content_block_delta", 0) + assertEventCount(t, events, "content_block_stop", 0) +} + +// TestProcessStreamLine_ReasoningBeatsContent verifies that when a single delta +// carries both "reasoning" and "content", BOTH are emitted rather than dropping +// the content. Some OpenAI-compatible adapters coalesce reasoning + content into +// one chunk; the old early-return behaviour silently lost the content. +// +// Correct output: thinking block (opened and closed) followed by text block, with +// the thinking block stop emitted before the text block start. +func TestProcessStreamLine_ReasoningBeatsContent(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Single chunk with both reasoning and content populated. + stream := mockOpenAIStream([]string{ + `data: {"id":"chatcmpl-rc","model":"deepseek-r1","choices":[{"delta":{"reasoning":"think","content":"text"},"index":0}]}` + "\n\n", + finishChunk("chatcmpl-rc", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // Both the thinking and the text must be present. + assert.Contains(t, body, `"thinking":"think"`) + assert.Contains(t, body, `"text":"text"`) + + // Two content_block_deltas: one thinking_delta, one text_delta. + assertEventCount(t, events, "content_block_delta", 2) + + deltas := findEventsByType(events, "content_block_delta") + require.Len(t, deltas, 2) + + firstDelta, ok := deltas[0]["delta"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "thinking_delta", firstDelta["type"], "first delta must be thinking_delta") + + secondDelta, ok := deltas[1]["delta"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "text_delta", secondDelta["type"], "second delta must be text_delta") + + // Two blocks (thinking + text), both properly closed. + assertContentBlockCount(t, events, 2) + assertBlocksClosed(t, events) + + // Thinking block must be fully stopped before the text block opens. + assertThinkingBlockTransitionOrder(t, events) +} + +// TestProcessStreamLine_ReasoningContentFieldName verifies the alternative field +// name "reasoning_content" (vLLM / SGLang / DeepSeek convention) is treated +// identically to "reasoning" and produces a thinking_delta. +func TestProcessStreamLine_ReasoningContentFieldName(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + `data: {"id":"chatcmpl-rcf","model":"qwq-32b","choices":[{"delta":{"reasoning_content":"deep think"},"index":0}]}` + "\n\n", + finishChunk("chatcmpl-rcf", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + assert.Contains(t, body, `"thinking":"deep think"`) + + // exactly one thinking block opened and closed + assertContentBlockCount(t, events, 1) + assertBlocksClosed(t, events) + + starts := findEventsByType(events, "content_block_start") + bt, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeThinking, bt) +} + +// TestProcessStreamLine_ToolCallContinuationNoID verifies the two-chunk tool +// call sequence: first chunk carries id+name (index=0), second carries only +// index+partial-args (no id, no name). The assembled args must appear in +// input_json_delta events and both deltas must reference block index 0. +// +// The args string must use the \\\" escaping convention (same as other streaming +// tests) because toolArgsChunk embeds the args into a JSON string value: raw +// braces would break the outer JSON and cause the chunk to be skipped. +func TestProcessStreamLine_ToolCallContinuationNoID(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + // first chunk: id + name establish the block + toolStartChunk("chatcmpl-cont", 0, "call_cont", "cont_tool"), + // second chunk: args only, no id or name; escaped per the test convention + toolArgsChunk("chatcmpl-cont", 0, `{\\\"x\\\":1}`), + finishChunk("chatcmpl-cont", "tool_calls"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // at least one input_json_delta must be present for the args chunk + assert.Contains(t, body, `"partial_json":`) + + // all input_json_delta events must reference block 0 + deltas := findEventsByType(events, "content_block_delta") + require.NotEmpty(t, deltas, "expected at least one delta for the args chunk") + for _, d := range deltas { + idx, ok := getContentBlockIndex(d) + require.True(t, ok) + assert.Equal(t, 0, idx, "all deltas must reference block 0") + } + + assertBlocksClosed(t, events) +} + +// TestProcessStreamLine_ToolCallAbsentIndexDefaultsToZero verifies that a tool +// call delta with no "index" key at all is treated as index 0. The JSON struct +// field uses *int so absent maps to nil, which defaults to 0. +func TestProcessStreamLine_ToolCallAbsentIndexDefaultsToZero(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Manually crafted raw SSE line — no "index" field in the tool_calls entry. + // toolArgsChunk always includes index, so we must build this by hand. + rawLine := "data: {\"id\":\"chatcmpl-noidx\",\"model\":\"claude-3-5-sonnet-20241022\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"id\":\"call_noidx\",\"type\":\"function\",\"function\":{\"name\":\"no_idx_tool\",\"arguments\":\"{\\\"x\\\":1}\"}}]},\"index\":0}]}\n\n" + + stream := mockOpenAIStream([]string{ + rawLine, + finishChunk("chatcmpl-noidx", "tool_calls"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // tool block must have been created (id + name present in that chunk) + assertToolPresent(t, body, "call_noidx", "no_idx_tool") + + // args appear as partial_json in the SSE frame — check for the field name + assert.Contains(t, body, `"partial_json":`) + deltas := findEventsByType(events, "content_block_delta") + require.NotEmpty(t, deltas, "at least one input_json_delta expected") + for _, d := range deltas { + idx, ok := getContentBlockIndex(d) + require.True(t, ok) + assert.Equal(t, 0, idx, "absent index must default to 0") + } +} + +// TestProcessStreamLine_MalformedChunkSkipped verifies that a malformed JSON +// chunk sandwiched between two valid text chunks is skipped without aborting the +// stream, and both valid chunks' text appears in the output. +func TestProcessStreamLine_MalformedChunkSkipped(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-bad2", "claude-3-5-sonnet-20241022", "before"), + "data: {not valid json at all}\n\n", + textChunk("chatcmpl-bad2", "", "after"), + finishChunk("chatcmpl-bad2", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + + // must not return an error + err := tr.TransformStreamingResponse(context.Background(), stream, recorder, nil) + require.NoError(t, err) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // both valid chunks must appear + assert.Contains(t, body, `"text":"before"`) + assert.Contains(t, body, `"text":"after"`) + + // stream must complete normally + assertHasEventType(t, events, "message_start") + assertHasEventType(t, events, "message_stop") +} + +// TestProcessStreamLine_UsageFinalChunk verifies that usage tokens in the final +// finish chunk are propagated correctly to the message_delta usage fields. +// The typed path reads PromptTokens/CompletionTokens as int directly, avoiding +// the float64 intermediate that the map-based path required. +func TestProcessStreamLine_UsageFinalChunk(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-usage2", "claude-3-5-sonnet-20241022", "hello"), + finishChunkWithUsage("chatcmpl-usage2", "stop", 15, 8), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + // assertUsageTokens checks message_delta usage + assertUsageTokens(t, events, 15, 8) + + // confirm the stop reason is mapped correctly + assertStopReason(t, events, "end_turn") + + // blocks must be balanced + assertBlocksClosed(t, events) + + // Confirm the values via the raw event too, since assertUsageTokens uses float64 + // (JSON numbers always unmarshal to float64 via map[string]interface{}). + messageDelta := getEventByType(events, "message_delta") + require.NotNil(t, messageDelta) + usage, ok := messageDelta["usage"].(map[string]interface{}) + require.True(t, ok, "message_delta must have usage") + + inputTokens, hasInput := usage["input_tokens"].(float64) + require.True(t, hasInput) + assert.Equal(t, float64(15), inputTokens) + + outputTokens, hasOutput := usage["output_tokens"].(float64) + require.True(t, hasOutput) + assert.Equal(t, float64(8), outputTokens) +} + +// TestProcessStreamLine_EmptyToolCallsArrayNoBlocks verifies that an empty +// tool_calls array ([]openAIToolCall with zero elements) does not open any +// content blocks. The range loop over an empty slice is a no-op. +func TestProcessStreamLine_EmptyToolCallsArrayNoBlocks(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-emtc", "claude-3-5-sonnet-20241022", "hello"), + // explicit empty tool_calls array + `data: {"id":"chatcmpl-emtc","choices":[{"delta":{"tool_calls":[]},"index":0}]}` + "\n\n", + finishChunk("chatcmpl-emtc", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + // only 1 block (the text block), no tool blocks + assertContentBlockCount(t, events, 1) + assertBlocksClosed(t, events) + + starts := findEventsByType(events, "content_block_start") + bt, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeText, bt) +} + +// TestProcessStreamLine_FinishReasonNullVsAbsent verifies that a delta where +// finish_reason is explicitly null does not update lastFinishReason and the +// stream stops with the default "end_turn" mapping. +func TestProcessStreamLine_FinishReasonNullVsAbsent(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-frnull", "claude-3-5-sonnet-20241022", "hello"), + // finish_reason explicitly null — must not override lastFinishReason + `data: {"id":"chatcmpl-frnull","choices":[{"delta":{},"index":0,"finish_reason":null}]}` + "\n\n", + finishChunk("chatcmpl-frnull", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + assertStopReason(t, events, "end_turn") +} + +// TestProcessStreamLine_ReasoningAndToolCallsInOneChunk verifies that when a single +// delta carries both "reasoning" and "tool_calls", both are processed. The thinking +// block must be fully stopped before the tool_use block opens. +func TestProcessStreamLine_ReasoningAndToolCallsInOneChunk(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Manually craft a chunk that includes both reasoning and tool_calls in one delta. + // This simulates a backend that coalesces chain-of-thought with a tool invocation. + combinedChunk := `data: {"id":"chatcmpl-rtc","model":"deepseek-r1","choices":[{"delta":{"reasoning":"I will call a tool","tool_calls":[{"index":0,"id":"call_rtc","type":"function","function":{"name":"search","arguments":""}}]},"index":0}]}` + "\n\n" + + stream := mockOpenAIStream([]string{ + combinedChunk, + toolArgsChunk("chatcmpl-rtc", 0, `{\\\"q\\\":\\\"test\\\"}`), + finishChunk("chatcmpl-rtc", "tool_calls"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // Both thinking content and the tool must be present. + assert.Contains(t, body, `"thinking":"I will call a tool"`) + assertToolPresent(t, body, "call_rtc", "search") + + // Two blocks: thinking and tool_use, both properly closed. + assertContentBlockCount(t, events, 2) + assertBlocksClosed(t, events) + + // Thinking block must be stopped before the tool_use block opens. + assertThinkingBlockTransitionOrder(t, events) + + starts := findEventsByType(events, "content_block_start") + require.Len(t, starts, 2) + thinkingType, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeThinking, thinkingType, "first block must be thinking") + toolType, _ := getContentBlockType(starts[1]) + assert.Equal(t, contentTypeToolUse, toolType, "second block must be tool_use") +} + +// TestProcessStreamLine_ReasoningContentAndToolCallsInOneChunk verifies that a +// single delta carrying reasoning, content, AND tool_calls together emits all +// three as separate, properly sequenced blocks. This exercises the three-way +// coalescing path: thinking block (start, thinking_delta, stop), then text block +// (start, text_delta, stop), then tool_use block (start, stop). Each block must +// be fully closed before the next one opens. +func TestProcessStreamLine_ReasoningContentAndToolCallsInOneChunk(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Single chunk carrying all three delta fields simultaneously. + // This simulates a backend that coalesces chain-of-thought, text, and a tool + // invocation into one SSE frame. + combinedChunk := `data: {"id":"chatcmpl-rct","model":"deepseek-r1","choices":[{"delta":{"reasoning":"think step","content":"answer text","tool_calls":[{"index":0,"id":"call_rct","type":"function","function":{"name":"lookup","arguments":""}}]},"index":0}]}` + "\n\n" + + stream := mockOpenAIStream([]string{ + combinedChunk, + toolArgsChunk("chatcmpl-rct", 0, `{\\\"q\\\":\\\"go\\\"}`), + finishChunk("chatcmpl-rct", "tool_calls"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // All three content types must be present in the body. + assert.Contains(t, body, `"thinking":"think step"`, "thinking content must be present") + assert.Contains(t, body, `"text":"answer text"`, "text content must be present") + assertToolPresent(t, body, "call_rct", "lookup") + + // Three blocks opened: thinking, text, tool_use — all closed. + assertContentBlockCount(t, events, 3) + assertBlocksClosed(t, events) + + // Verify block order: thinking first, then text, then tool_use. + starts := findEventsByType(events, "content_block_start") + require.Len(t, starts, 3, "expected exactly 3 content_block_start events") + + thinkingType, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeThinking, thinkingType, "first block must be thinking") + + textType, _ := getContentBlockType(starts[1]) + assert.Equal(t, contentTypeText, textType, "second block must be text") + + toolType, _ := getContentBlockType(starts[2]) + assert.Equal(t, contentTypeToolUse, toolType, "third block must be tool_use") + + // Thinking block must be fully stopped before the text block opens. + assertThinkingBlockTransitionOrder(t, events) + + // Text block (index 1) must be stopped before the tool_use block (index 2) opens. + stops := findEventsByType(events, "content_block_stop") + require.GreaterOrEqual(t, len(stops), 2, "at least thinking and text blocks must be stopped before tool_use opens") + textStopIdx := -1 + toolStartIdx := -1 + for i, ev := range events { + switch ev["_event_type"] { + case "content_block_stop": + if idx, ok := getContentBlockIndex(ev); ok && idx == 1 && textStopIdx == -1 { + textStopIdx = i + } + case "content_block_start": + bt, ok := getContentBlockType(ev) + if ok && bt == contentTypeToolUse && toolStartIdx == -1 { + toolStartIdx = i + } + } + } + if textStopIdx > -1 && toolStartIdx > -1 { + assert.Less(t, textStopIdx, toolStartIdx, "text block must stop before tool_use block opens") + } +} + +// validateStreamingParseTests is a compile-time guard — if any of the above +// helper dependencies shift, this will fail to compile rather than silently pass. +func validateStreamingParseTests() { + _ = strings.Builder{} +} diff --git a/internal/adapter/translator/anthropic/streaming_test.go b/internal/adapter/translator/anthropic/streaming_test.go index aadecc8a..e65e0561 100644 --- a/internal/adapter/translator/anthropic/streaming_test.go +++ b/internal/adapter/translator/anthropic/streaming_test.go @@ -4,11 +4,13 @@ import ( "context" "fmt" "io" + "net/http/httptest" "strings" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/thushan/olla/internal/core/constants" ) func TestTransformStreamingResponse_SimpleText(t *testing.T) { @@ -452,11 +454,14 @@ func TestTransformStreamingResponse_MultipleToolsSequential(t *testing.T) { func TestTransformStreamingResponse_InterleavedToolArguments(t *testing.T) { translator := newTestTranslator() - // simulate arguments arriving interleaved between tools (common in streaming) + // Interleaved tool starts followed by interleaved args. Under the single-active-block + // invariant, toolStart(1) closes block 0 before opening block 1. Args for tool 0 + // arriving after its block was stopped trigger an SSE error event and abort the stream; + // corrupt tool input is worse than a clean failure the client can retry. stream := mockOpenAIStream([]string{ toolStartChunk("chatcmpl-int", 0, "call_A", "toolA"), toolStartChunk("chatcmpl-int", 1, "call_B", "toolB"), - // interleave chunks + // args for tool 0 arrive after its block was stopped by tool 1's start toolArgsChunk("chatcmpl-int", 0, `{\\\"data\\\"`), toolArgsChunk("chatcmpl-int", 1, `{\\\"value\\\"`), toolArgsChunk("chatcmpl-int", 0, `:123}`), @@ -465,13 +470,141 @@ func TestTransformStreamingResponse_InterleavedToolArguments(t *testing.T) { doneChunk(), }) - recorder := executeTransform(t, translator, stream) + recorder := httptest.NewRecorder() + err := translator.TransformStreamingResponse(context.Background(), stream, recorder, nil) + // stream must be aborted with the sentinel error + require.ErrorIs(t, err, errInterleavedToolArguments) + body := recorder.Body.String() - // verify both tools present with complete arguments despite interleaving + // Both tool starts must have been emitted before the abort. assertToolPresent(t, body, "call_A", "toolA") assertToolPresent(t, body, "call_B", "toolB") - assertContainsAll(t, body, []string{`123`, `456`}) + + // SSE error event must be on the wire. + assert.Contains(t, body, "event: error") + assert.Contains(t, body, `"type":"error"`) + assert.Contains(t, body, `"type":"api_error"`) + + // message_delta and message_stop must NOT appear after the error event. + assert.NotContains(t, body, "event: message_delta") + assert.NotContains(t, body, "event: message_stop") +} + +// TestTransformStreamingResponse_PreInitBufferedArgsFlush verifies that args +// arriving before the id+name chunk are emitted as the first input_json_delta +// immediately after content_block_start, making the pre-init case lossless. +func TestTransformStreamingResponse_PreInitBufferedArgsFlush(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Construct a stream where the args chunk for tool 0 arrives before the + // toolStart (id+name) chunk. This simulates a backend that delivers function + // arguments in the first delta before the function metadata chunk. + preInitArgs := `{\\\"prefix\\\"` + laterArgs := `:\\\"value\\\"}` + + // Build the raw SSE bytes manually so we can put the args chunk first. + stream := mockOpenAIStream([]string{ + toolArgsChunk("chatcmpl-preinit", 0, preInitArgs), // args arrive before id+name + toolStartChunk("chatcmpl-preinit", 0, "call_pre", "pre_tool"), + toolArgsChunk("chatcmpl-preinit", 0, laterArgs), + finishChunk("chatcmpl-preinit", "tool_calls"), + doneChunk(), + }) + + recorder := executeTransform(t, tr, stream) + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // Tool block must be present. + assertToolPresent(t, body, "call_pre", "pre_tool") + + // Both the pre-init and later args must appear in input_json_delta events. + assert.Contains(t, body, "prefix") + assert.Contains(t, body, "value") + + // There must be at least two input_json_delta events: one for the pre-init + // flush and one (or more) for the later args. + deltas := findEventsByType(events, "content_block_delta") + require.GreaterOrEqual(t, len(deltas), 2, "expected at least 2 deltas: pre-init flush + later args") + + // All deltas must reference the single tool block (index 0). + for _, d := range deltas { + idx, ok := getContentBlockIndex(d) + require.True(t, ok) + assert.Equal(t, 0, idx, "all deltas should reference block 0") + } + + assertBlocksClosed(t, events) +} + +// TestTransformStreamingResponse_MessageStartCacheFields verifies that message_start +// carries the cache_creation_input_tokens and cache_read_input_tokens fields (both 0). +func TestTransformStreamingResponse_MessageStartCacheFields(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-cache", "claude-3-5-sonnet-20241022", "Hi"), + finishChunk("chatcmpl-cache", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + start := getEventByType(events, "message_start") + require.NotNil(t, start) + message, ok := start["message"].(map[string]interface{}) + require.True(t, ok) + usage, ok := message["usage"].(map[string]interface{}) + require.True(t, ok) + + _, hasCacheCreate := usage["cache_creation_input_tokens"] + _, hasCacheRead := usage["cache_read_input_tokens"] + assert.True(t, hasCacheCreate, "message_start usage must include cache_creation_input_tokens") + assert.True(t, hasCacheRead, "message_start usage must include cache_read_input_tokens") + assert.Equal(t, float64(0), usage["cache_creation_input_tokens"]) + assert.Equal(t, float64(0), usage["cache_read_input_tokens"]) + + // stop_reason and stop_sequence must be present in the message object + _, hasStopReason := message["stop_reason"] + _, hasStopSeq := message["stop_sequence"] + assert.True(t, hasStopReason, "message_start message must include stop_reason") + assert.True(t, hasStopSeq, "message_start message must include stop_sequence") +} + +// TestTransformStreamingResponse_MessageDeltaCacheFields verifies that message_delta +// carries the cache_creation_input_tokens and cache_read_input_tokens fields (both 0). +func TestTransformStreamingResponse_MessageDeltaCacheFields(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-deltacache", "claude-3-5-sonnet-20241022", "Hi"), + finishChunkWithUsage("chatcmpl-deltacache", "stop", 5, 3), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + delta := getEventByType(events, "message_delta") + require.NotNil(t, delta) + usage, ok := delta["usage"].(map[string]interface{}) + require.True(t, ok) + + _, hasCacheCreate := usage["cache_creation_input_tokens"] + _, hasCacheRead := usage["cache_read_input_tokens"] + assert.True(t, hasCacheCreate, "message_delta usage must include cache_creation_input_tokens") + assert.True(t, hasCacheRead, "message_delta usage must include cache_read_input_tokens") + assert.Equal(t, float64(0), usage["cache_creation_input_tokens"]) + assert.Equal(t, float64(0), usage["cache_read_input_tokens"]) } func TestTransformStreamingResponse_ToolTextToolTransitions(t *testing.T) { @@ -669,16 +802,14 @@ func TestTransformStreamingResponse_EmptyToolCallsArray(t *testing.T) { func TestTransformStreamingResponse_ManyTools(t *testing.T) { translator := newTestTranslator() - // create stream with 15 tools + // create stream with 15 tools, sending each start then its own args immediately + // (sequential, not interleaved) so no block is already stopped when args arrive. chunks := []string{ textChunk("chatcmpl-many", "claude-3-5-sonnet-20241022", "Processing many tools"), } - // add 15 tools with interleaved arguments for i := range 15 { chunks = append(chunks, toolStartChunk("chatcmpl-many", i, fmt.Sprintf("call_%d", i), fmt.Sprintf("tool_%d", i))) - } - for i := range 15 { chunks = append(chunks, toolArgsChunk("chatcmpl-many", i, fmt.Sprintf(`{\\\"n\\\":%d}`, i))) } @@ -730,3 +861,494 @@ func TestTransformStreamingResponse_RapidBlockTypeSwitch(t *testing.T) { assertToolPresent(t, body, "call_2", "tool2") assertTextContent(t, body, "Text3") } + +// TestTransformStreamingResponse_MessageStartInputTokensSeeded verifies that when +// the handler injects a pre-computed token estimate into context (via +// ContextInputTokensKey), the message_start event carries that non-zero value. +// +// vLLM and lmdeploy both populate real input_tokens in message_start from the first +// upstream chunk; Olla's translation path seeds the value from the request body +// estimator to match that behaviour without buffering the stream. +// +// output_tokens in message_start must remain 0 per the Anthropic spec. +func TestTransformStreamingResponse_MessageStartInputTokensSeeded(t *testing.T) { + trans := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-seeded", "claude-3-5-sonnet-20241022", "Hello"), + finishChunk("chatcmpl-seeded", "stop"), + doneChunk(), + }) + + // Simulate the handler injecting a pre-computed estimate into context. + const seedTokens = 42 + ctx := context.WithValue(context.Background(), constants.ContextInputTokensKey, seedTokens) + + recorder := httptest.NewRecorder() + err := trans.TransformStreamingResponse(ctx, stream, recorder, nil) + require.NoError(t, err) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + messageStart := getEventByType(events, "message_start") + require.NotNil(t, messageStart, "message_start event must be present") + + message, ok := messageStart["message"].(map[string]interface{}) + require.True(t, ok, "message_start must have a message field") + + usage, ok := message["usage"].(map[string]interface{}) + require.True(t, ok, "message_start.message must have a usage field") + + inputTokens, ok := usage["input_tokens"].(float64) + require.True(t, ok, "usage must have input_tokens") + assert.Equal(t, float64(seedTokens), inputTokens, + "message_start input_tokens must reflect the seeded estimate") + + outputTokens, ok := usage["output_tokens"].(float64) + require.True(t, ok, "usage must have output_tokens") + assert.Equal(t, float64(0), outputTokens, + "message_start output_tokens must be 0 per Anthropic spec") +} + +// TestTransformStreamingResponse_MessageStartNoSeedIsZero confirms baseline behaviour: +// without a seeded estimate in context, message_start input_tokens is 0. +// This covers the passthrough path and any caller that does not inject the key. +func TestTransformStreamingResponse_MessageStartNoSeedIsZero(t *testing.T) { + trans := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-noseed", "claude-3-5-sonnet-20241022", "Hi"), + finishChunk("chatcmpl-noseed", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + err := trans.TransformStreamingResponse(context.Background(), stream, recorder, nil) + require.NoError(t, err) + + events := parseAnthropicEvents(t, recorder.Body.String()) + messageStart := getEventByType(events, "message_start") + require.NotNil(t, messageStart) + + message := messageStart["message"].(map[string]interface{}) + usage := message["usage"].(map[string]interface{}) + inputTokens := usage["input_tokens"].(float64) + assert.Equal(t, float64(0), inputTokens, "no seed in context should give 0 input_tokens in message_start") +} + +// TestTransformStreamingResponse_TwoSequentialToolCallsExactSequence pins the full SSE +// event order for a stream containing two back-to-back tool calls. Every +// content_block_start must be paired with a content_block_stop and indices must +// correspond to the correct block throughout. +func TestTransformStreamingResponse_TwoSequentialToolCallsExactSequence(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + toolStartChunk("chatcmpl-seq2", 0, "call_seq_0", "tool_alpha"), + toolArgsChunk("chatcmpl-seq2", 0, `{`), + toolArgsChunk("chatcmpl-seq2", 0, `\\\"x\\\":1}`), + toolStartChunk("chatcmpl-seq2", 1, "call_seq_1", "tool_beta"), + toolArgsChunk("chatcmpl-seq2", 1, `{`), + toolArgsChunk("chatcmpl-seq2", 1, `\\\"y\\\":2}`), + finishChunk("chatcmpl-seq2", "tool_calls"), + doneChunk(), + }) + + recorder := executeTransform(t, tr, stream) + events := parseAnthropicEvents(t, recorder.Body.String()) + + // Exact event sequence required by the Anthropic SSE spec. + assertEventSequence(t, events, []string{ + "message_start", + "content_block_start", // index 0, tool_alpha + "content_block_delta", // index 0, partial_json { + "content_block_delta", // index 0, partial_json \"x\":1} + "content_block_stop", // index 0 + "content_block_start", // index 1, tool_beta + "content_block_delta", // index 1, partial_json { + "content_block_delta", // index 1, partial_json \"y\":2} + "content_block_stop", // index 1 + "message_delta", + "message_stop", + }) + + starts := findEventsByType(events, "content_block_start") + stops := findEventsByType(events, "content_block_stop") + require.Len(t, starts, 2, "two tool blocks must be started") + require.Len(t, stops, 2, "both tool blocks must be stopped") + + // Block 0 is tool_alpha at index 0. + idx0, _ := getContentBlockIndex(starts[0]) + assert.Equal(t, 0, idx0) + bt0, _ := getContentBlockType(starts[0]) + assert.Equal(t, contentTypeToolUse, bt0) + + // Block 1 is tool_beta at index 1. + idx1, _ := getContentBlockIndex(starts[1]) + assert.Equal(t, 1, idx1) + bt1, _ := getContentBlockType(starts[1]) + assert.Equal(t, contentTypeToolUse, bt1) + + // Each stop event carries the correct index. + stopIdx0, _ := getContentBlockIndex(stops[0]) + assert.Equal(t, 0, stopIdx0) + stopIdx1, _ := getContentBlockIndex(stops[1]) + assert.Equal(t, 1, stopIdx1) + + // Deltas must carry the owning block's index, not the current block at emit time. + deltas := findEventsByType(events, "content_block_delta") + require.Len(t, deltas, 4) + di0, _ := getContentBlockIndex(deltas[0]) + assert.Equal(t, 0, di0, "first delta must be for block 0") + di1, _ := getContentBlockIndex(deltas[1]) + assert.Equal(t, 0, di1, "second delta must be for block 0") + di2, _ := getContentBlockIndex(deltas[2]) + assert.Equal(t, 1, di2, "third delta must be for block 1") + di3, _ := getContentBlockIndex(deltas[3]) + assert.Equal(t, 1, di3, "fourth delta must be for block 1") + + assertStopReason(t, events, "tool_use") +} + +// TestTransformStreamingResponse_TextThenToolExactSequence pins the event order when +// a text block transitions into a tool call. The text block must be stopped before +// the tool block starts; this is already guarded by assertBlockTransitionOrder but +// this test also verifies the exact full sequence. +func TestTransformStreamingResponse_TextThenToolExactSequence(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-txttool", "claude-3-5-sonnet-20241022", "Thinking..."), + toolStartChunk("chatcmpl-txttool", 0, "call_after_text", "lookup"), + toolArgsChunk("chatcmpl-txttool", 0, `{\\\"k\\\":\\\"v\\\"}`), + finishChunk("chatcmpl-txttool", "tool_calls"), + doneChunk(), + }) + + recorder := executeTransform(t, tr, stream) + events := parseAnthropicEvents(t, recorder.Body.String()) + + assertEventSequence(t, events, []string{ + "message_start", + "content_block_start", // index 0, text + "content_block_delta", // text_delta "Thinking..." + "content_block_stop", // index 0 stopped before tool starts + "content_block_start", // index 1, tool_use + "content_block_delta", // input_json_delta + "content_block_stop", // index 1 + "message_delta", + "message_stop", + }) + + assertBlocksClosed(t, events) + assertBlockTransitionOrder(t, events) + + // text block is index 0, tool block is index 1 + starts := findEventsByType(events, "content_block_start") + textIdx, _ := getContentBlockIndex(starts[0]) + assert.Equal(t, 0, textIdx) + toolIdx, _ := getContentBlockIndex(starts[1]) + assert.Equal(t, 1, toolIdx) +} + +// TestTransformStreamingResponse_LateArgDeltaDropped verifies that a late delta +// arriving for a tool whose block has already been stopped aborts the stream with +// errInterleavedToolArguments and emits a spec-valid SSE error event. The stream +// must not emit message_delta or message_stop after the error. +func TestTransformStreamingResponse_LateArgDeltaDropped(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // tool 0 starts and gets an arg, then tool 1 starts (which closes tool 0), + // then a late arg arrives for tool 0 -- must abort the stream. + stream := mockOpenAIStream([]string{ + toolStartChunk("chatcmpl-late", 0, "call_early", "first_tool"), + toolArgsChunk("chatcmpl-late", 0, `{\\\"a\\\":1}`), + toolStartChunk("chatcmpl-late", 1, "call_late_target", "second_tool"), + toolArgsChunk("chatcmpl-late", 0, `{\\\"late\\\":true}`), // late for tool 0 -- aborts + toolArgsChunk("chatcmpl-late", 1, `{\\\"b\\\":2}`), + finishChunk("chatcmpl-late", "tool_calls"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + err := tr.TransformStreamingResponse(context.Background(), stream, recorder, nil) + require.ErrorIs(t, err, errInterleavedToolArguments) + + body := recorder.Body.String() + + // SSE error event must be present and no finalise events must follow. + assert.Contains(t, body, "event: error") + assert.Contains(t, body, `"type":"api_error"`) + assert.NotContains(t, body, "event: message_delta") + assert.NotContains(t, body, "event: message_stop") + + // Both tool block starts (emitted before the abort) must be present. + assertToolPresent(t, body, "call_early", "first_tool") + assertToolPresent(t, body, "call_late_target", "second_tool") +} + +// TestTransformStreamingResponse_UsageOnlyFinalChunk verifies that a terminal chunk +// with choices:[] and a populated usage object -- the format emitted by vLLM and +// other backends that honour include_usage=true -- is correctly captured and reflected +// in the message_delta usage. Previously the empty-choices guard returned early before +// reading chunk.Usage, so message_delta kept the pre-stream estimate or showed +// output_tokens:0. +func TestTransformStreamingResponse_UsageOnlyFinalChunk(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Stream layout: + // 1. text chunk (carries model name) + // 2. finish chunk (no usage, just the finish reason) + // 3. usage-only chunk with choices:[] (the backend's include_usage terminal chunk) + // 4. [DONE] + // + // The usage in step 3 must win over any estimate seeded earlier. + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-uonly", "claude-3-5-sonnet-20241022", "Hi"), + finishChunk("chatcmpl-uonly", "stop"), + usageOnlyChunk("chatcmpl-uonly", 37, 14), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + // Real token counts must appear in message_delta, not 0 or a stale estimate. + assertUsageTokens(t, events, 37, 14) +} + +// TestTransformStreamingResponse_UsageOnlyChunkOverridesEstimate verifies that the +// real backend usage in a terminal choices:[] chunk overwrites a non-zero estimate +// that was seeded into context before the stream started. This guards against the +// common case where the estimate is close but not exact. +func TestTransformStreamingResponse_UsageOnlyChunkOverridesEstimate(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-override", "claude-3-5-sonnet-20241022", "Hello"), + finishChunk("chatcmpl-override", "stop"), + usageOnlyChunk("chatcmpl-override", 55, 20), + doneChunk(), + }) + + // Seed a non-zero estimate that differs from the real usage. + const seedTokens = 10 + ctx := context.WithValue(context.Background(), constants.ContextInputTokensKey, seedTokens) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(ctx, stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + // Real counts from the usage-only chunk must replace the seeded estimate. + assertUsageTokens(t, events, 55, 20) +} + +// TestTransformStreamingResponse_RepeatedToolIDNameDoesNotReinit verifies that when +// a backend repeats the tool id + name on continuation chunks (alongside args), only +// a single content_block_start is emitted for that tool. Re-initialising would close +// the active block and open a duplicate, replaying already-buffered args. +func TestTransformStreamingResponse_RepeatedToolIDNameDoesNotReinit(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // First chunk: id + name + partial args (some backends send all three together). + // Second chunk: id + name AGAIN + more args (the bug trigger). + // Third chunk: args only (normal continuation). + firstChunk := `data: {"id":"chatcmpl-repeat","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_rep","type":"function","function":{"name":"rep_tool","arguments":"{\\\"a\\\":"}}]},"index":0}]}` + "\n\n" + repeatChunk := `data: {"id":"chatcmpl-repeat","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_rep","type":"function","function":{"name":"rep_tool","arguments":"1,"}}]},"index":0}]}` + "\n\n" + finalChunk := `data: {"id":"chatcmpl-repeat","choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\\\"b\\\":2}"}}]},"index":0}]}` + "\n\n" + + stream := mockOpenAIStream([]string{ + firstChunk, + repeatChunk, + finalChunk, + finishChunk("chatcmpl-repeat", "tool_calls"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + body := recorder.Body.String() + events := parseAnthropicEvents(t, body) + + // Exactly one tool block must have been started, not two. + assertContentBlockCount(t, events, 1) + assertBlocksClosed(t, events) + + // The single tool block must carry the correct id and name. + assertToolPresent(t, body, "call_rep", "rep_tool") + + // All input_json_delta events must reference block index 0 - no spurious block 1. + deltas := findEventsByType(events, "content_block_delta") + require.NotEmpty(t, deltas, "expected at least one input_json_delta") + for _, d := range deltas { + idx, ok := getContentBlockIndex(d) + require.True(t, ok) + assert.Equal(t, 0, idx, "all deltas must reference block 0") + } + + // Three arg chunks must have produced at least three input_json_delta events. + // This confirms args from all chunks were emitted, not dropped on the repeated id+name. + require.GreaterOrEqual(t, len(deltas), 3, "expected one delta per arg chunk") +} + +// TestTransformStreamingResponse_SynthesisedOutputTokensWhenNoBackendUsage verifies +// that when the backend emits no usage in its stream, finalizeStream synthesises a +// non-zero output_tokens estimate from the accumulated content length (chars/4, min 1). +// This prevents Anthropic clients (including Claude Code) from receiving output_tokens:0. +func TestTransformStreamingResponse_SynthesisedOutputTokensWhenNoBackendUsage(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Stream with several content chunks but no usage chunk at all - the common case + // for most OpenAI-compatible backends that do not support include_usage. + content := "Hello, world! This is synthesised." // 34 chars -> 34/4 = 8 + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-synth", "claude-3-5-sonnet-20241022", content), + finishChunk("chatcmpl-synth", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + delta := getEventByType(events, "message_delta") + require.NotNil(t, delta, "message_delta must be present") + + usage, ok := delta["usage"].(map[string]interface{}) + require.True(t, ok, "message_delta must have usage") + + outputTokens, ok := usage["output_tokens"].(float64) + require.True(t, ok, "usage must have output_tokens") + + // Must be non-zero and approximately match chars/4 for the accumulated content. + assert.Greater(t, outputTokens, float64(0), + "output_tokens must be synthesised to a non-zero value when backend sends no usage") + assert.InDelta(t, float64(len(content)/4), outputTokens, 2, + "synthesised output_tokens should roughly match chars/4 of the accumulated content") +} + +// TestTransformStreamingResponse_RealUsageWinsOverSynthesis verifies that when the +// backend does provide usage via a terminal choices:[] chunk, the real completion_tokens +// value is used in message_delta rather than the chars/4 estimate. +func TestTransformStreamingResponse_RealUsageWinsOverSynthesis(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + const realCompletionTokens = 20 + + // The usageOnlyChunk emits a choices:[] terminal chunk -- the same format used by + // vLLM and the ollamock terminal usage chunk. The translator must write 20, not the + // estimate derived from the accumulated content. + stream := mockOpenAIStream([]string{ + textChunk("chatcmpl-realusage", "claude-3-5-sonnet-20241022", "Hello world"), + finishChunk("chatcmpl-realusage", "stop"), + usageOnlyChunk("chatcmpl-realusage", 10, realCompletionTokens), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + delta := getEventByType(events, "message_delta") + require.NotNil(t, delta, "message_delta must be present") + + usage, ok := delta["usage"].(map[string]interface{}) + require.True(t, ok, "message_delta must have usage") + + outputTokens, ok := usage["output_tokens"].(float64) + require.True(t, ok, "usage must have output_tokens") + + // Real backend value must win; synthesis must not fire when state.outputTokens != 0. + assert.Equal(t, float64(realCompletionTokens), outputTokens, + "real backend output_tokens must take precedence over the chars/4 estimate") +} + +// TestTransformStreamingResponse_SynthesisedOutputTokensEmptyContent verifies that a +// stream with no text or thinking content (e.g. a tool-only response with no visible +// text) synthesises output_tokens:0, which is correct -- the client-visible content +// length is genuinely zero and a min-1 clamp would be misleading. +// The min-1 clamp is only applied when charCount > 0. +func TestTransformStreamingResponse_SynthesisedOutputTokensEmptyContent(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Only a finish chunk, no content or usage -- simulates the truly-nothing case. + stream := mockOpenAIStream([]string{ + finishChunk("chatcmpl-empty-synth", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + delta := getEventByType(events, "message_delta") + require.NotNil(t, delta, "message_delta must be present") + + usage, ok := delta["usage"].(map[string]interface{}) + require.True(t, ok, "message_delta must have usage") + + outputTokens, ok := usage["output_tokens"].(float64) + require.True(t, ok, "usage must have output_tokens") + + // No content, no real usage -> output_tokens stays 0 (no spurious min-1 clamp). + assert.Equal(t, float64(0), outputTokens, + "output_tokens must be 0 when there is no content and no backend usage") +} + +// TestTransformStreamingResponse_SynthesisedOutputTokensThinkingContent verifies that +// thinking (reasoning) block text is included in the synthesis estimate alongside +// regular text content, ensuring chain-of-thought responses get a plausible token count. +func TestTransformStreamingResponse_SynthesisedOutputTokensThinkingContent(t *testing.T) { + t.Parallel() + tr := newTestTranslator() + + // Reasoning chunk followed by a short text chunk, no backend usage. + thinking := "Let me reason through this carefully." + visible := "The answer is 42." + stream := mockOpenAIStream([]string{ + reasoningChunk("chatcmpl-think-synth", "claude-3-5-sonnet-20241022", thinking), + textChunk("chatcmpl-think-synth", "", visible), + finishChunk("chatcmpl-think-synth", "stop"), + doneChunk(), + }) + + recorder := httptest.NewRecorder() + require.NoError(t, tr.TransformStreamingResponse(context.Background(), stream, recorder, nil)) + + events := parseAnthropicEvents(t, recorder.Body.String()) + + delta := getEventByType(events, "message_delta") + require.NotNil(t, delta, "message_delta must be present") + + usage, ok := delta["usage"].(map[string]interface{}) + require.True(t, ok, "message_delta must have usage") + + outputTokens, ok := usage["output_tokens"].(float64) + require.True(t, ok, "usage must have output_tokens") + + // Both thinking and visible content contribute to the estimate, so the token + // count must be at least 1 and greater than visible text alone. + totalChars := len(thinking) + len(visible) + minExpected := float64(1) + maxExpected := float64(totalChars/4 + 2) // +2 tolerance for integer division + assert.GreaterOrEqual(t, outputTokens, minExpected, + "synthesised output_tokens must be >= 1 when thinking+text content exists") + assert.LessOrEqual(t, outputTokens, maxExpected, + "synthesised output_tokens must not far exceed (thinking+text chars)/4") +} diff --git a/internal/adapter/translator/anthropic/streaming_test_helpers.go b/internal/adapter/translator/anthropic/streaming_test_helpers.go index 2e29ab36..22524f67 100644 --- a/internal/adapter/translator/anthropic/streaming_test_helpers.go +++ b/internal/adapter/translator/anthropic/streaming_test_helpers.go @@ -24,7 +24,11 @@ func newTestTranslator() *Translator { loggerCfg := &logger.Config{Level: "error", Theme: "default"} log, _, _ := logger.New(loggerCfg) styledLog := logger.NewPlainStyledLogger(log) - return NewTranslator(styledLog, createStreamingTestConfig()) + tr, err := NewTranslator(styledLog, createStreamingTestConfig()) + if err != nil { + panic("newTestTranslator: " + err.Error()) + } + return tr } // createStreamingTestConfig creates a minimal config for streaming tests. @@ -92,6 +96,40 @@ func doneChunk() string { return "data: [DONE]\n\n" } +// usageOnlyChunk creates an OpenAI SSE chunk with choices:[] and a populated usage +// object. This is the format vLLM (and other backends that respect include_usage=true) +// use to deliver final token counts as a separate terminal chunk rather than attaching +// them to the finish chunk. The message ID is included for completeness; clients +// correlate by stream position rather than by ID. +func usageOnlyChunk(messageID string, promptTokens, completionTokens int) string { + return fmt.Sprintf( + "data: {\"id\":\"%s\",\"choices\":[],\"usage\":{\"prompt_tokens\":%d,\"completion_tokens\":%d,\"total_tokens\":%d}}\n\n", + messageID, promptTokens, completionTokens, promptTokens+completionTokens, + ) +} + +// reasoningChunk creates an OpenAI SSE chunk carrying reasoning text using the +// "reasoning" field name (Ollama / LM Studio / Lemonade convention). +func reasoningChunk(messageID, model, reasoning string) string { + if model != "" { + return fmt.Sprintf("data: {\"id\":\"%s\",\"model\":\"%s\",\"choices\":[{\"delta\":{\"reasoning\":\"%s\"},\"index\":0}]}\n\n", + messageID, model, reasoning) + } + return fmt.Sprintf("data: {\"id\":\"%s\",\"choices\":[{\"delta\":{\"reasoning\":\"%s\"},\"index\":0}]}\n\n", + messageID, reasoning) +} + +// reasoningContentChunk creates an OpenAI SSE chunk carrying reasoning text using the +// "reasoning_content" field name (vLLM / SGLang / DeepSeek convention). +func reasoningContentChunk(messageID, model, reasoning string) string { + if model != "" { + return fmt.Sprintf("data: {\"id\":\"%s\",\"model\":\"%s\",\"choices\":[{\"delta\":{\"reasoning_content\":\"%s\"},\"index\":0}]}\n\n", + messageID, model, reasoning) + } + return fmt.Sprintf("data: {\"id\":\"%s\",\"choices\":[{\"delta\":{\"reasoning_content\":\"%s\"},\"index\":0}]}\n\n", + messageID, reasoning) +} + // modelChunk creates the first chunk with model information. // Use this as the first chunk in a stream when you need to specify the model. func modelChunk(messageID, model, content string) string { @@ -397,6 +435,42 @@ func assertUsageTokens(t *testing.T, events []map[string]interface{}, expectedIn assert.Equal(t, float64(expectedOutput), output, "output_tokens should match") } +// assertThinkingContent validates that thinking text appears in a thinking_delta event. +func assertThinkingContent(t *testing.T, body string, thinking string) { + t.Helper() + assert.Contains(t, body, fmt.Sprintf(`"thinking":"%s"`, thinking)) +} + +// assertThinkingBlockTransitionOrder validates that a thinking block stops before any +// text or tool block starts. Anthropic requires thinking blocks to precede content. +func assertThinkingBlockTransitionOrder(t *testing.T, events []map[string]interface{}) { + t.Helper() + + thinkingStopIdx := -1 + firstContentOrToolStartIdx := -1 + + for i, event := range events { + switch event["_event_type"] { + case "content_block_stop": + // Find the stop for the thinking block (always opened first, so index 0). + if idx, ok := getContentBlockIndex(event); ok && idx == 0 && thinkingStopIdx == -1 { + // Confirm it was actually the thinking block by checking the corresponding start. + thinkingStopIdx = i + } + case "content_block_start": + bt, ok := getContentBlockType(event) + if ok && bt != contentTypeThinking && firstContentOrToolStartIdx == -1 { + firstContentOrToolStartIdx = i + } + } + } + + if thinkingStopIdx > -1 && firstContentOrToolStartIdx > -1 { + assert.Less(t, thinkingStopIdx, firstContentOrToolStartIdx, + "thinking block must be stopped before any text/tool block starts") + } +} + // assertBlockTransitionOrder validates that text block stops before tool block starts. // This is critical for proper content block lifecycle when transitioning types. func assertBlockTransitionOrder(t *testing.T, events []map[string]interface{}) { diff --git a/internal/adapter/translator/anthropic/token_count.go b/internal/adapter/translator/anthropic/token_count.go index fd192265..46ec483c 100644 --- a/internal/adapter/translator/anthropic/token_count.go +++ b/internal/adapter/translator/anthropic/token_count.go @@ -11,6 +11,27 @@ import ( "github.com/thushan/olla/internal/adapter/translator" ) +// countTokensWireResponse is the Anthropic-conformant wire shape for +// POST /v1/messages/count_tokens. The spec returns {"input_tokens":N} only; +// output_tokens and total_tokens are not part of this endpoint's contract. +// Confirmed against vLLM, lmdeploy, bifrost, and litellm reference impls. +type countTokensWireResponse struct { + InputTokens int `json:"input_tokens"` +} + +// EstimateInputTokens implements translator.InputTokenEstimator. +// Called by the streaming path to seed input_tokens in message_start before the +// upstream usage chunk arrives (which only comes at the end of the stream). +// vLLM and lmdeploy both populate real input_tokens in message_start; this brings +// Olla's translation path in line with that behaviour. +func (t *Translator) EstimateInputTokens(originalBodyBytes []byte) int { + var req AnthropicRequest + if err := json.Unmarshal(originalBodyBytes, &req); err != nil { + return 0 + } + return estimateTokensFromRequest(&req) +} + // token estimation for claude code compatibility func (t *Translator) CountTokens(ctx context.Context, r *http.Request) (*translator.TokenCountResponse, error) { // bounded read to prevent OOM attacks @@ -42,12 +63,18 @@ func (t *Translator) CountTokens(ctx context.Context, r *http.Request) (*transla tokenCount := estimateTokensFromRequest(&req) return &translator.TokenCountResponse{ - InputTokens: tokenCount, - OutputTokens: 0, // zero output tokens for count endpoint - TotalTokens: tokenCount, + InputTokens: tokenCount, }, nil } +// SerialiseCountTokens implements translator.TokenCountSerializer. +// Returns only {"input_tokens":N} — the Anthropic spec for this endpoint does not +// include output_tokens or total_tokens. Confirmed against vLLM, lmdeploy, bifrost, +// and litellm reference implementations. +func (t *Translator) SerialiseCountTokens(resp *translator.TokenCountResponse) ([]byte, error) { + return json.Marshal(countTokensWireResponse{InputTokens: resp.InputTokens}) +} + // character-based token estimation // simple algorithm from python reference: chars / 4 func estimateTokensFromRequest(req *AnthropicRequest) int { @@ -62,6 +89,12 @@ func estimateTokensFromRequest(req *AnthropicRequest) int { totalChars += countMessageChars(&msg) } + // Claude Code sends 10-20k tokens of tool schemas per request; omitting tools + // causes a large systematic undercount that misseeds message_start input_tokens. + for i := range req.Tools { + totalChars += countToolDefinitionChars(&req.Tools[i]) + } + tokenCount := totalChars / 4 if tokenCount < 1 { tokenCount = 1 @@ -70,6 +103,21 @@ func estimateTokensFromRequest(req *AnthropicRequest) int { return tokenCount } +// countToolDefinitionChars estimates character cost of a single tool definition. +// Counts name + description + JSON-marshalled input schema, matching the chars/4 heuristic. +func countToolDefinitionChars(tool *AnthropicTool) int { + if tool == nil { + return 0 + } + n := len(tool.Name) + len(tool.Description) + if tool.InputSchema != nil { + if schemaJSON, err := json.Marshal(tool.InputSchema); err == nil { + n += len(schemaJSON) + } + } + return n +} + // system prompt char counting // handles string and content block formats func countSystemChars(system interface{}) int { diff --git a/internal/adapter/translator/anthropic/token_count_test.go b/internal/adapter/translator/anthropic/token_count_test.go index 70e65ed7..d263ab29 100644 --- a/internal/adapter/translator/anthropic/token_count_test.go +++ b/internal/adapter/translator/anthropic/token_count_test.go @@ -7,10 +7,13 @@ import ( "io" "net/http" "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestCountTokens(t *testing.T) { - trans := NewTranslator(createTestLogger(), createTestConfig()) + trans := mustNewTranslator(createTestLogger(), createTestConfig()) tests := []struct { name string @@ -223,16 +226,6 @@ func TestCountTokens(t *testing.T) { t.Fatal("Expected non-nil response") } - // Check output tokens is always 0 - if resp.OutputTokens != 0 { - t.Errorf("Expected OutputTokens=0, got %d", resp.OutputTokens) - } - - // Check input tokens matches total tokens - if resp.InputTokens != resp.TotalTokens { - t.Errorf("Expected InputTokens=%d to match TotalTokens=%d", resp.InputTokens, resp.TotalTokens) - } - // Verify token count if tt.expectedExact > 0 { if resp.InputTokens != tt.expectedExact { @@ -245,14 +238,13 @@ func TestCountTokens(t *testing.T) { } } - t.Logf("Token count: %d (input=%d, output=%d, total=%d)", - resp.InputTokens, resp.InputTokens, resp.OutputTokens, resp.TotalTokens) + t.Logf("Token count: input=%d", resp.InputTokens) }) } } func TestCountTokensWithRawJSON(t *testing.T) { - trans := NewTranslator(createTestLogger(), createTestConfig()) + trans := mustNewTranslator(createTestLogger(), createTestConfig()) tests := []struct { name string @@ -322,7 +314,7 @@ func TestCountTokensWithRawJSON(t *testing.T) { } func TestCountTokensErrors(t *testing.T) { - trans := NewTranslator(createTestLogger(), createTestConfig()) + trans := mustNewTranslator(createTestLogger(), createTestConfig()) tests := []struct { name string @@ -380,7 +372,7 @@ func TestCountTokensErrors(t *testing.T) { // TestCountTokensMatchesPythonReference verifies our implementation matches // the Python reference from anthropic-proxy.py func TestCountTokensMatchesPythonReference(t *testing.T) { - trans := NewTranslator(createTestLogger(), createTestConfig()) + trans := mustNewTranslator(createTestLogger(), createTestConfig()) // This test case directly mirrors the Python implementation logic testCase := AnthropicRequest{ @@ -429,8 +421,95 @@ func TestCountTokensMatchesPythonReference(t *testing.T) { t.Errorf("Expected %d tokens to match Python reference, got %d", expectedTokens, resp.InputTokens) } } + +// TestCountTokensWireResponse verifies that SerialiseCountTokens emits only +// {"input_tokens":N} — no output_tokens or total_tokens fields. +// Anthropic's spec, vLLM, lmdeploy, bifrost, and litellm all define this shape. +func TestCountTokensWireResponse(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + + reqBody := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "Hello world"}] + }`) + + req, err := http.NewRequest("POST", "/v1/messages/count_tokens", bytes.NewReader(reqBody)) + if err != nil { + t.Fatalf("failed to create request: %v", err) + } + + resp, err := trans.CountTokens(context.Background(), req) + if err != nil { + t.Fatalf("CountTokens failed: %v", err) + } + if resp.InputTokens == 0 { + t.Fatal("expected non-zero InputTokens") + } + + wireBytes, err := trans.SerialiseCountTokens(resp) + if err != nil { + t.Fatalf("SerialiseCountTokens failed: %v", err) + } + + var wireMap map[string]interface{} + if err := json.Unmarshal(wireBytes, &wireMap); err != nil { + t.Fatalf("wire response is not valid JSON: %v", err) + } + + // Must contain input_tokens + if _, ok := wireMap["input_tokens"]; !ok { + t.Error("wire response missing input_tokens field") + } + + // Must NOT contain output_tokens or total_tokens + if _, ok := wireMap["output_tokens"]; ok { + t.Error("wire response must not contain output_tokens (not part of Anthropic count_tokens spec)") + } + if _, ok := wireMap["total_tokens"]; ok { + t.Error("wire response must not contain total_tokens (not part of Anthropic count_tokens spec)") + } + + // Must contain exactly one field + if len(wireMap) != 1 { + t.Errorf("wire response must have exactly 1 field, got %d: %v", len(wireMap), wireMap) + } +} + +// TestEstimateInputTokens verifies the streaming token seeding path. +func TestEstimateInputTokens(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + + body := []byte(`{ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "Hello, this is a test message for token estimation."}] + }`) + + estimate := trans.EstimateInputTokens(body) + if estimate <= 0 { + t.Errorf("expected positive token estimate for non-empty prompt, got %d", estimate) + } +} + +// TestEstimateInputTokens_InvalidBody returns 0 without panicking on malformed input. +func TestEstimateInputTokens_InvalidBody(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + + estimate := trans.EstimateInputTokens([]byte(`{not json`)) + if estimate != 0 { + t.Errorf("expected 0 for invalid body, got %d", estimate) + } +} + func BenchmarkCountTokens(b *testing.B) { - trans := NewTranslator(createTestLogger(), createTestConfig()) + trans := mustNewTranslator(createTestLogger(), createTestConfig()) reqBody := []byte(`{ "model": "claude-3-5-sonnet-20241022", @@ -447,8 +526,86 @@ func BenchmarkCountTokens(b *testing.B) { } } +// TestEstimateInputTokens_WithTools verifies that tool definitions are included in the +// character estimate. Claude Code sends large tool schemas (10-20k tokens each), so +// omitting them causes a large systematic undercount. +func TestEstimateInputTokens_WithTools(t *testing.T) { + t.Parallel() + + baseReq := AnthropicRequest{ + Model: "claude-3-5-sonnet-20241022", + MaxTokens: 1024, + Messages: []AnthropicMessage{ + {Role: "user", Content: "What is the weather?"}, + }, + } + + reqWithTools := baseReq + reqWithTools.Tools = []AnthropicTool{ + { + Name: "get_weather", + Description: "Retrieves current weather for a given location.", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "location": map[string]interface{}{ + "type": "string", + "description": "City and state, e.g. Sydney, NSW", + }, + }, + "required": []string{"location"}, + }, + }, + } + + baseCount := estimateTokensFromRequest(&baseReq) + withToolsCount := estimateTokensFromRequest(&reqWithTools) + + assert.Greater(t, withToolsCount, baseCount, + "request with tools must yield a larger token estimate than the same request without") +} + +// TestEstimateInputTokens_ToolOnlyDeltaMatchesCharsDiv4 verifies that the tool schema +// contribution to the estimate is roughly len(schema JSON) / 4. +func TestEstimateInputTokens_ToolOnlyDeltaMatchesCharsDiv4(t *testing.T) { + t.Parallel() + + schema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{"type": "string"}, + }, + } + tool := AnthropicTool{ + Name: "search", + Description: "Search the web", + InputSchema: schema, + } + + // Manually compute expected char count for the tool. + schemaJSON, err := json.Marshal(schema) + require.NoError(t, err) + expectedChars := len(tool.Name) + len(tool.Description) + len(schemaJSON) + expectedTokens := expectedChars / 4 + + got := countToolDefinitionChars(&tool) + assert.Equal(t, expectedChars, got, "countToolDefinitionChars must match name+description+schema len") + + // estimateTokensFromRequest for a tools-only request (empty messages, single tool) + // must be at least expectedTokens. + req := AnthropicRequest{ + Model: "test", + MaxTokens: 1, + Messages: []AnthropicMessage{{Role: "user", Content: "hi"}}, + Tools: []AnthropicTool{tool}, + } + total := estimateTokensFromRequest(&req) + assert.GreaterOrEqual(t, total, expectedTokens, + "total estimate must be at least the tool's contribution") +} + func BenchmarkCountTokensLargeRequest(b *testing.B) { - trans := NewTranslator(createTestLogger(), createTestConfig()) + trans := mustNewTranslator(createTestLogger(), createTestConfig()) // Large request with multiple messages and content blocks messages := make([]map[string]interface{}, 50) diff --git a/internal/adapter/translator/anthropic/translator.go b/internal/adapter/translator/anthropic/translator.go index bf69de24..cdd7e925 100644 --- a/internal/adapter/translator/anthropic/translator.go +++ b/internal/adapter/translator/anthropic/translator.go @@ -25,19 +25,20 @@ type Translator struct { maxMessageSize int64 // derived from config } -// NewTranslator creates a new Anthropic translator instance -// Uses a buffer pool to reduce GC pressure during high-throughput operations -// Accepts configuration for request size limits and streaming behaviour -func NewTranslator(log logger.StyledLogger, cfg config.AnthropicTranslatorConfig) *Translator { - // Create buffer pool with 4KB initial capacity - // This size fits most chat completions without reallocation +// NewTranslator creates a new Anthropic translator instance. +// Uses a buffer pool to reduce GC pressure during high-throughput operations. +// Accepts configuration for request size limits and streaming behaviour. +// Returns an error if the buffer pool cannot be constructed (unreachable in +// production because bytes.NewBuffer never returns nil, but exposed here to +// honour the no-panic policy). +func NewTranslator(log logger.StyledLogger, cfg config.AnthropicTranslatorConfig) (*Translator, error) { + // Create buffer pool with 4KB initial capacity. + // This size fits most chat completions without reallocation. bufferPool, err := pool.NewLitePool(func() *bytes.Buffer { return bytes.NewBuffer(make([]byte, 0, 4096)) }) if err != nil { - // This should never happen as the constructor is validated - log.Error("Failed to create buffer pool", "error", err) - panic("translator: failed to initialise buffer pool") + return nil, fmt.Errorf("translator: failed to initialise buffer pool: %w", err) } // Apply defaults if needed @@ -61,7 +62,7 @@ func NewTranslator(log logger.StyledLogger, cfg config.AnthropicTranslatorConfig config: cfg, maxMessageSize: maxSize, inspector: insp, - } + }, nil } // Name returns the translator identifier @@ -117,7 +118,10 @@ func (t *Translator) WriteError(w http.ResponseWriter, err error, statusCode int "error", err.Error(), "status", statusCode) - // Map HTTP status codes to Anthropic error types + // Map HTTP status codes to Anthropic error types. + // overloaded_error is reserved for Anthropic's own infrastructure being overloaded; + // a generic backend 503 is an api_error per litellm and bifrost conventions. + // timeout_error maps to 504/408 as defined by litellm and the Anthropic error taxonomy. errorType := "api_error" switch statusCode { case http.StatusBadRequest: @@ -128,10 +132,13 @@ func (t *Translator) WriteError(w http.ResponseWriter, err error, statusCode int errorType = "permission_error" case http.StatusNotFound: errorType = "not_found_error" + case http.StatusRequestEntityTooLarge: + // Anthropic error taxonomy: oversized request body maps to request_too_large. + errorType = "request_too_large" case http.StatusTooManyRequests: errorType = "rate_limit_error" - case http.StatusServiceUnavailable: - errorType = "overloaded_error" + case http.StatusRequestTimeout, http.StatusGatewayTimeout: + errorType = "timeout_error" } // Anthropic error format @@ -143,6 +150,14 @@ func (t *Translator) WriteError(w http.ResponseWriter, err error, statusCode int }, } + // Best-effort: propagate the Olla request ID into the Anthropic error envelope. + // SDKs read the request-id response header and the top-level request_id body field + // for correlation. Only set when the header was already written by the handler. + if reqID := w.Header().Get(constants.HeaderXOllaRequestID); reqID != "" { + w.Header().Set("request-id", reqID) + errorResp["request_id"] = reqID + } + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) w.WriteHeader(statusCode) diff --git a/internal/adapter/translator/anthropic/translator_test.go b/internal/adapter/translator/anthropic/translator_test.go index 6cac9f34..04f6566b 100644 --- a/internal/adapter/translator/anthropic/translator_test.go +++ b/internal/adapter/translator/anthropic/translator_test.go @@ -14,7 +14,7 @@ import ( // TestGetAPIPath tests the PathProvider interface implementation func TestGetAPIPath(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) path := translator.GetAPIPath() assert.Equal(t, "/olla/anthropic/v1/messages", path, "should return the correct Anthropic API path") @@ -65,10 +65,13 @@ func TestWriteError(t *testing.T) { expectedStatus: http.StatusTooManyRequests, }, { + // 503 from a backend means the backend is unavailable, not that Anthropic's + // own infra is overloaded. overloaded_error is an Anthropic-origin signal; + // a generic gateway 503 maps to api_error per litellm/bifrost conventions. name: "service_unavailable", err: errors.New("service overloaded"), statusCode: http.StatusServiceUnavailable, - expectedType: "overloaded_error", + expectedType: "api_error", expectedStatus: http.StatusServiceUnavailable, }, { @@ -82,7 +85,7 @@ func TestWriteError(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) rec := httptest.NewRecorder() translator.WriteError(rec, tt.err, tt.statusCode) @@ -113,7 +116,7 @@ func TestWriteError(t *testing.T) { // TestWriteError_ErrorFormat tests Anthropic error format compliance func TestWriteError_ErrorFormat(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) rec := httptest.NewRecorder() testErr := errors.New("test error message") @@ -135,14 +138,28 @@ func TestWriteError_ErrorFormat(t *testing.T) { // TestName tests the Name method func TestName(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) assert.Equal(t, "anthropic", translator.Name()) } +// TestNewTranslator_Success verifies the happy path returns a non-nil translator +// and no error, confirming the error-return refactor did not break normal construction. +func TestNewTranslator_Success(t *testing.T) { + t.Parallel() + + tr, err := NewTranslator(createTestLogger(), createTestConfig()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tr == nil { + t.Fatal("expected non-nil translator") + } +} + // TestWriteError_JSONEncodingFailure tests handling of JSON encoding errors // This test is mostly for coverage, as encoding errors are rare func TestWriteError_JSONEncodingSuccess(t *testing.T) { - translator := NewTranslator(createTestLogger(), createTestConfig()) + translator := mustNewTranslator(createTestLogger(), createTestConfig()) rec := httptest.NewRecorder() // Standard error that should encode successfully @@ -156,3 +173,129 @@ func TestWriteError_JSONEncodingSuccess(t *testing.T) { err := json.Unmarshal(rec.Body.Bytes(), &response) assert.NoError(t, err, "response should be valid JSON") } + +// TestWriteError_ConformanceMapping verifies Anthropic-conformant error type mappings. +// +// Key correctness constraints: +// - 503 must map to api_error, NOT overloaded_error. overloaded_error is an Anthropic-origin +// signal (litellm maps it only for responses from Anthropic's own infrastructure). A backend +// 503 means the upstream gateway is unavailable — that is api_error territory. +// - 504 and 408 must map to timeout_error per litellm's error taxonomy. +// - Response body must be the Anthropic error envelope: {"type":"error","error":{"type":"...","message":"..."}}. +func TestWriteError_ConformanceMapping(t *testing.T) { + t.Parallel() + + cases := []struct { + status int + expectedType string + }{ + {http.StatusBadRequest, "invalid_request_error"}, + {http.StatusUnauthorized, "authentication_error"}, + {http.StatusForbidden, "permission_error"}, + {http.StatusNotFound, "not_found_error"}, + {http.StatusRequestEntityTooLarge, "request_too_large"}, + {http.StatusTooManyRequests, "rate_limit_error"}, + {http.StatusInternalServerError, "api_error"}, + {http.StatusBadGateway, "api_error"}, + // 503 from a backend is NOT Anthropic-origin overloaded; it is api_error. + {http.StatusServiceUnavailable, "api_error"}, + {http.StatusGatewayTimeout, "timeout_error"}, + {http.StatusRequestTimeout, "timeout_error"}, + } + + for _, tc := range cases { + + t.Run(http.StatusText(tc.status), func(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + rec := httptest.NewRecorder() + trans.WriteError(rec, errors.New("test error"), tc.status) + + require.Equal(t, tc.status, rec.Code) + require.Equal(t, constants.ContentTypeJSON, rec.Header().Get(constants.HeaderContentType)) + + var body map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + + // Anthropic error envelope: {"type":"error","error":{...}} + assert.Equal(t, "error", body["type"], "outer type must be 'error'") + + errObj, ok := body["error"].(map[string]interface{}) + require.True(t, ok, "body.error must be an object") + assert.Equal(t, tc.expectedType, errObj["type"], + "status %d should map to %s", tc.status, tc.expectedType) + assert.NotEmpty(t, errObj["message"], "error.message must not be empty") + }) + } +} + +// TestWriteError_413RequestTooLarge verifies that an oversized body (413) maps to the +// Anthropic "request_too_large" error type, not the generic "api_error". +func TestWriteError_413RequestTooLarge(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + rec := httptest.NewRecorder() + + trans.WriteError(rec, errors.New("request body exceeds maximum size"), http.StatusRequestEntityTooLarge) + + require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code) + require.Equal(t, constants.ContentTypeJSON, rec.Header().Get(constants.HeaderContentType)) + + var body map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + + assert.Equal(t, "error", body["type"]) + + errObj, ok := body["error"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, "request_too_large", errObj["type"], + "413 must map to request_too_large per the Anthropic error taxonomy") + assert.NotEmpty(t, errObj["message"]) +} + +// TestWriteError_RequestIDPresent verifies that when X-Olla-Request-ID is already set on +// the response, WriteError copies it to the request-id response header and to the top-level +// request_id field of the error body. Anthropic SDKs read both for correlation. +func TestWriteError_RequestIDPresent(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + rec := httptest.NewRecorder() + + const reqID = "olla-req-abc123" + rec.Header().Set(constants.HeaderXOllaRequestID, reqID) + + trans.WriteError(rec, errors.New("something failed"), http.StatusInternalServerError) + + require.Equal(t, http.StatusInternalServerError, rec.Code) + + // request-id response header must be set + assert.Equal(t, reqID, rec.Header().Get("request-id")) + + // request_id must appear as a top-level field in the JSON body + var body map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + assert.Equal(t, reqID, body["request_id"], + "request_id top-level field must match the Olla request ID header") +} + +// TestWriteError_RequestIDAbsent verifies that when X-Olla-Request-ID is not set, +// WriteError does not fabricate a request_id and does not set the request-id header. +func TestWriteError_RequestIDAbsent(t *testing.T) { + t.Parallel() + + trans := mustNewTranslator(createTestLogger(), createTestConfig()) + rec := httptest.NewRecorder() + + // Do NOT set X-Olla-Request-ID. + trans.WriteError(rec, errors.New("something failed"), http.StatusBadRequest) + + assert.Empty(t, rec.Header().Get("request-id"), "request-id header must not be set when source header is absent") + + var body map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + _, hasRequestID := body["request_id"] + assert.False(t, hasRequestID, "request_id must not appear in body when source header is absent") +} diff --git a/internal/adapter/translator/anthropic/types.go b/internal/adapter/translator/anthropic/types.go index b4141495..9fd7308c 100644 --- a/internal/adapter/translator/anthropic/types.go +++ b/internal/adapter/translator/anthropic/types.go @@ -60,13 +60,14 @@ type AnthropicMessage struct { } // ContentBlock represents different types of content in messages -// Anthropic uses a block-based content model for text, images, tool use, and tool results +// Anthropic uses a block-based content model for text, images, tool use, tool results, and thinking. type ContentBlock struct { Content interface{} `json:"content,omitempty"` // for tool_result Source *ImageSource `json:"source,omitempty"` Input map[string]interface{} `json:"input,omitempty"` - Type string `json:"type"` // "text", "image", "tool_use", "tool_result" + Type string `json:"type"` // "text", "image", "tool_use", "tool_result", "thinking" Text string `json:"text,omitempty"` + Thinking string `json:"thinking,omitempty"` // populated for type="thinking" response blocks ID string `json:"id,omitempty"` Name string `json:"name,omitempty"` ToolUseID string `json:"tool_use_id,omitempty"` @@ -168,6 +169,12 @@ type InputJSONDelta struct { PartialJSON string `json:"partial_json"` } +// ThinkingDelta represents incremental reasoning/thinking text +type ThinkingDelta struct { + Type string `json:"type"` // "thinking_delta" + Thinking string `json:"thinking"` +} + // ContentBlockStop marks the end of a content block type ContentBlockStop struct { Type string `json:"type"` // "content_block_stop" diff --git a/internal/adapter/translator/types.go b/internal/adapter/translator/types.go index be14c879..81b0c527 100644 --- a/internal/adapter/translator/types.go +++ b/internal/adapter/translator/types.go @@ -56,6 +56,21 @@ type TokenCounter interface { CountTokens(ctx context.Context, r *http.Request) (*TokenCountResponse, error) } +// optional interface for translators that need to control the wire serialisation +// of their token count response. When implemented, the handler uses SerialiseCountTokens +// instead of encoding TokenCountResponse directly, allowing translators to emit only +// the fields their spec defines (e.g. Anthropic: {"input_tokens":N} only). +type TokenCountSerializer interface { + SerialiseCountTokens(resp *TokenCountResponse) ([]byte, error) +} + +// optional interface for translators that can cheaply estimate input token count +// from a pre-buffered original request body. Used to seed input_tokens in streaming +// message_start events before the upstream usage chunk arrives. +type InputTokenEstimator interface { + EstimateInputTokens(originalBodyBytes []byte) int +} + // represents token count result type TokenCountResponse struct { InputTokens int `json:"input_tokens"` diff --git a/internal/adapter/unifier/lifecycle_unifier.go b/internal/adapter/unifier/lifecycle_unifier.go index 03a3bde0..b6db6ae7 100644 --- a/internal/adapter/unifier/lifecycle_unifier.go +++ b/internal/adapter/unifier/lifecycle_unifier.go @@ -24,9 +24,8 @@ type LifecycleUnifier struct { discoveryClient DiscoveryClient logger logger.StyledLogger - cleanupCtx context.Context - endpointManager *EndpointManager - stateTransitions chan domain.StateTransition + cleanupCtx context.Context + endpointManager *EndpointManager cleanupCancel context.CancelFunc config Config @@ -40,11 +39,10 @@ func NewLifecycleUnifier(config Config, logger logger.StyledLogger) ports.ModelU } return &LifecycleUnifier{ - unifier: NewDefaultUnifier(), - endpointManager: NewEndpointManager(config, logger), - config: config, - logger: logger, - stateTransitions: make(chan domain.StateTransition, 100), + unifier: NewDefaultUnifier(), + endpointManager: NewEndpointManager(config, logger), + config: config, + logger: logger, } } @@ -64,11 +62,6 @@ func (u *LifecycleUnifier) Start(ctx context.Context) error { go u.cleanupRoutine() } - if u.config.EnableStateTransitionLogging { - u.cleanupWg.Add(1) - go u.stateTransitionLogger() - } - u.logger.Info("Lifecycle unifier started", "model_ttl", u.config.ModelTTL, "cleanup_interval", u.config.CleanupInterval, @@ -102,7 +95,6 @@ func (u *LifecycleUnifier) Stop(ctx context.Context) error { return ctx.Err() } - close(u.stateTransitions) return nil } @@ -181,33 +173,6 @@ func (u *LifecycleUnifier) performCleanup() { u.endpointManager.CleanupOrphaned(activeEndpoints) } -func (u *LifecycleUnifier) stateTransitionLogger() { - defer u.cleanupWg.Done() - - for { - select { - case <-u.cleanupCtx.Done(): - return - case transition, ok := <-u.stateTransitions: - if !ok { - return - } - if transition.Error != nil { - u.logger.Warn("State transition", - "from", transition.From, - "to", transition.To, - "reason", transition.Reason, - "error", transition.Error) - } else { - u.logger.Debug("State transition", - "from", transition.From, - "to", transition.To, - "reason", transition.Reason) - } - } - } -} - func (u *LifecycleUnifier) UnifyModel(ctx context.Context, sourceModel *domain.ModelInfo, endpoint *domain.Endpoint) (*domain.UnifiedModel, error) { models := []*domain.ModelInfo{sourceModel} unified, err := u.UnifyModels(ctx, models, endpoint) diff --git a/internal/app/handlers/application.go b/internal/app/handlers/application.go index 5935c1b5..b004cee1 100644 --- a/internal/app/handlers/application.go +++ b/internal/app/handlers/application.go @@ -29,29 +29,43 @@ type SecurityAdapters struct { logger logger.StyledLogger } -// CreateChainMiddleware creates middleware that applies the full security chain with enhanced logging +// CreateChainMiddleware creates middleware that applies the full security chain with enhanced logging. +// When concrete security adapters are available (production path), we delegate to them so that +// per-validator status codes (429 rate-limit, 413 body-too-large) are preserved. The abstract +// securityChain path is kept as a fallback for test contexts where only the chain is wired. func (s *SecurityAdapters) CreateChainMiddleware() func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { - // Chain the middleware: logging -> access logging -> security -> handler + // Wrap with logging so every proxy request is recorded regardless of which + // security path runs below. withLogging := middleware.EnhancedLoggingMiddleware(s.logger)(next) withAccessLogging := middleware.AccessLoggingMiddleware(s.logger)(withLogging) + if s.securityAdapters != nil { + // Delegate to the concrete adapter chain. It sets the correct status codes + // (429 + Retry-After/X-RateLimit-* for rate limiting, 413 for oversized + // bodies) rather than flattening everything to 403. + concreteChain := s.securityAdapters.CreateChainMiddleware()(withAccessLogging) + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + concreteChain.ServeHTTP(w, r) + }) + } + + // Fallback: abstract chain only (e.g. unit tests that inject securityChain + // directly without wiring the full security.Adapters). return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if s.securityChain != nil { - // Create security request from HTTP request secReq := ports.SecurityRequest{ - ClientID: r.RemoteAddr, // This would normally be extracted better + ClientID: r.RemoteAddr, Endpoint: r.URL.Path, Method: r.Method, BodySize: r.ContentLength, - HeaderSize: 0, // Would need to calculate + HeaderSize: 0, Headers: r.Header, IsHealthCheck: r.URL.Path == "/internal/health", } result, err := s.securityChain.Validate(r.Context(), secReq) if err != nil || !result.Allowed { - // Write appropriate error response http.Error(w, "Security validation failed", http.StatusForbidden) return } @@ -169,7 +183,10 @@ func NewApplication( return nil, fmt.Errorf("invalid Anthropic translator config: %w", err) } - anthropicTranslator := anthropic.NewTranslator(logger, cfg.Translators.Anthropic) + anthropicTranslator, err := anthropic.NewTranslator(logger, cfg.Translators.Anthropic) + if err != nil { + return nil, fmt.Errorf("failed to create Anthropic translator: %w", err) + } translatorRegistry.Register("anthropic", anthropicTranslator) logger.Info("Registered Anthropic translator", diff --git a/internal/app/handlers/handler_security_status_codes_test.go b/internal/app/handlers/handler_security_status_codes_test.go new file mode 100644 index 00000000..5bec7326 --- /dev/null +++ b/internal/app/handlers/handler_security_status_codes_test.go @@ -0,0 +1,202 @@ +package handlers + +// Regression tests for the security-middleware status-code bug: +// Before the fix, proxy routes received 403 "Security validation failed" for +// both rate-limit rejections and oversized-body rejections, because +// SecurityAdapters.CreateChainMiddleware called the abstract securityChain.Validate +// and hard-coded http.StatusForbidden for any rejection. +// +// The correct behaviour (enforced here): +// - Rate-limited request → 429 + Retry-After + X-RateLimit-* headers +// - Oversized body → 413 + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/thushan/olla/internal/adapter/security" + "github.com/thushan/olla/internal/config" + "github.com/thushan/olla/internal/core/ports" +) + +// buildTightSecurityAdapters returns a *security.Adapters whose rate limiter +// immediately rejects any request (1 req/min, burst 0) and whose size +// validator rejects bodies > 10 bytes. +func buildTightSecurityAdapters(t *testing.T) *security.Adapters { + t.Helper() + + cfg := &config.Config{ + Server: config.ServerConfig{ + RateLimits: config.ServerRateLimits{ + GlobalRequestsPerMinute: 1, + PerIPRequestsPerMinute: 1, + BurstSize: 0, + }, + RequestLimits: config.ServerRequestLimits{ + MaxBodySize: 10, + MaxHeaderSize: 0, + }, + }, + } + _, adapters := security.NewSecurityServices(cfg, &mockStatsCollector{}, &mockStyledLogger{}) + t.Cleanup(adapters.Stop) + return adapters +} + +// buildSecurityAdaptersWithSizeOnly returns adapters with no rate limiting but a +// strict 10-byte body cap, to isolate the 413 path. +func buildSecurityAdaptersWithSizeOnly(t *testing.T) *security.Adapters { + t.Helper() + + cfg := &config.Config{ + Server: config.ServerConfig{ + RateLimits: config.ServerRateLimits{ + // Zero values mean the rate limiter allows everything. + GlobalRequestsPerMinute: 0, + PerIPRequestsPerMinute: 0, + BurstSize: 0, + }, + RequestLimits: config.ServerRequestLimits{ + MaxBodySize: 10, + MaxHeaderSize: 0, + }, + }, + } + _, adapters := security.NewSecurityServices(cfg, &mockStatsCollector{}, &mockStyledLogger{}) + t.Cleanup(adapters.Stop) + return adapters +} + +// okHandler is a trivial handler that records whether it was reached. +func okHandler(reached *bool) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + *reached = true + w.WriteHeader(http.StatusOK) + } +} + +// alwaysRejectValidator is a SecurityValidator that always denies requests. +type alwaysRejectValidator struct{} + +func (alwaysRejectValidator) Name() string { return "always-reject" } + +func (alwaysRejectValidator) Validate(_ context.Context, _ ports.SecurityRequest) (ports.SecurityResult, error) { + return ports.SecurityResult{Allowed: false, Reason: "test rejection"}, nil +} + +// TestSecurityChainMiddleware_RateLimit_Returns429 verifies that a proxy route +// gets 429 (not 403) when the rate limiter rejects the request, and that the +// Retry-After and X-RateLimit-* response headers are present. +// +// burst=0 means every reserve call returns !OK, so the very first request is rejected. +func TestSecurityChainMiddleware_RateLimit_Returns429(t *testing.T) { + t.Parallel() + + adapters := buildTightSecurityAdapters(t) + + sa := &SecurityAdapters{ + securityAdapters: adapters, + logger: &mockStyledLogger{}, + } + + reached := false + chain := sa.CreateChainMiddleware()(okHandler(&reached)) + + req := httptest.NewRequest(http.MethodPost, "/olla/proxy/v1/chat/completions", + strings.NewReader(`{"model":"llama3"}`)) + req.Header.Set("Content-Type", "application/json") + req.RemoteAddr = "1.2.3.4:9999" + + w := httptest.NewRecorder() + chain.ServeHTTP(w, req) + + if w.Code == http.StatusForbidden { + t.Fatal("rate limit returned 403 Forbidden — status code is being flattened by the abstract chain path") + } + if w.Code != http.StatusTooManyRequests { + t.Errorf("expected 429 for rate-limited request, got %d", w.Code) + } + + // The concrete rate-limit middleware sets these headers unconditionally. + if w.Header().Get("Retry-After") == "" { + t.Error("rate-limited response missing Retry-After header") + } + if w.Header().Get("X-RateLimit-Limit") == "" { + t.Error("rate-limited response missing X-RateLimit-Limit header") + } + if w.Header().Get("X-RateLimit-Remaining") == "" { + t.Error("rate-limited response missing X-RateLimit-Remaining header") + } + if reached { + t.Error("handler was reached; rate-limited request should have been rejected") + } +} + +// TestSecurityChainMiddleware_OversizedBody_Returns413 verifies that a proxy route +// gets 413 (not 403) when the request body exceeds the configured limit. +func TestSecurityChainMiddleware_OversizedBody_Returns413(t *testing.T) { + t.Parallel() + + adapters := buildSecurityAdaptersWithSizeOnly(t) + + sa := &SecurityAdapters{ + securityAdapters: adapters, + logger: &mockStyledLogger{}, + } + + reached := false + chain := sa.CreateChainMiddleware()(okHandler(&reached)) + + // Body is 50 bytes, well above the 10-byte cap. + body := strings.Repeat("x", 50) + req := httptest.NewRequest(http.MethodPost, "/olla/proxy/v1/chat/completions", + strings.NewReader(body)) + req.ContentLength = int64(len(body)) + req.Header.Set("Content-Type", "application/json") + req.RemoteAddr = "1.2.3.4:9999" + + w := httptest.NewRecorder() + chain.ServeHTTP(w, req) + + if w.Code == http.StatusForbidden { + t.Fatal("oversized body returned 403 Forbidden — status code is being flattened by the abstract chain path") + } + if w.Code != http.StatusRequestEntityTooLarge { + t.Errorf("expected 413 for oversized body, got %d", w.Code) + } + if reached { + t.Error("handler was reached; oversized request should have been rejected") + } +} + +// TestSecurityChainMiddleware_FallbackPath_Returns403 verifies that the abstract +// chain fallback (no securityAdapters wired) still returns 403 for rejected +// requests. This covers test contexts where only securityChain is available. +func TestSecurityChainMiddleware_FallbackPath_Returns403(t *testing.T) { + t.Parallel() + + // Wire only the abstract chain, no concrete securityAdapters. + rejectingChain := ports.NewSecurityChain(alwaysRejectValidator{}) + sa := &SecurityAdapters{ + securityChain: rejectingChain, + logger: &mockStyledLogger{}, + // securityAdapters intentionally nil — tests the fallback path. + } + + reached := false + chain := sa.CreateChainMiddleware()(okHandler(&reached)) + + req := httptest.NewRequest(http.MethodPost, "/olla/proxy/", strings.NewReader("hi")) + w := httptest.NewRecorder() + chain.ServeHTTP(w, req) + + if w.Code != http.StatusForbidden { + t.Errorf("abstract-chain fallback expected 403, got %d", w.Code) + } + if reached { + t.Error("handler should not be reached when abstract chain rejects") + } +} diff --git a/internal/app/handlers/handler_translation.go b/internal/app/handlers/handler_translation.go index e663d4a2..a3c9a8df 100644 --- a/internal/app/handlers/handler_translation.go +++ b/internal/app/handlers/handler_translation.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "sync" + "sync/atomic" "time" "github.com/thushan/olla/internal/adapter/translator" @@ -16,6 +17,7 @@ import ( "github.com/thushan/olla/internal/core/domain" "github.com/thushan/olla/internal/core/ports" "github.com/thushan/olla/internal/util" + "github.com/tidwall/gjson" ) // executePassthroughRequest handles requests that can be forwarded directly to backends @@ -44,12 +46,16 @@ func (a *Application) executePassthroughRequest( return } - // The proxy selector chooses the actual endpoint later, so only override the - // translator default when every candidate backend agrees on the same native - // path. Mixed native-Anthropic backends can require different paths (for - // example DMR uses /anthropic/v1/messages while vLLM uses /v1/messages). - if targetPath := a.resolvePassthroughTargetPath(endpoints, passthroughReq.TargetPath); targetPath != "" { - passthroughReq.TargetPath = targetPath + // The proxy selector chooses the actual endpoint later. When capable backends + // disagree on their native Anthropic path (e.g. DMR uses /anthropic/v1/messages + // while vLLM uses /v1/messages), we restrict to the largest path-compatible + // subset so the selected backend and the request path always agree. + resolvedPath, filteredEndpoints := a.resolvePassthroughTargetPath(endpoints, passthroughReq.TargetPath) + if resolvedPath != "" { + passthroughReq.TargetPath = resolvedPath + } + if len(filteredEndpoints) > 0 { + endpoints = filteredEndpoints } // Update proxy request details - capture streaming flag for accurate metrics @@ -94,33 +100,77 @@ func (a *Application) executePassthroughRequest( pr.stats.EndTime = time.Now() } -// resolvePassthroughTargetPath returns the native Anthropic path only when every -// candidate backend agrees on the same one. The proxy selector picks the real -// endpoint after this runs, so a mixed fleet (e.g. DMR on /anthropic/v1/messages -// alongside vLLM on /v1/messages) can't be given a single correct path; in that -// case we keep the translator's neutral default rather than risk a 404 on whichever -// backend is chosen. -func (a *Application) resolvePassthroughTargetPath(endpoints []*domain.Endpoint, defaultPath string) string { +// resolvePassthroughTargetPath resolves the native Anthropic messages path and +// the endpoint subset that must be used with it. +// +// Uniform fleets: all candidates share the same MessagesPath. Returns that path +// and nil (caller keeps original endpoint list unchanged). +// +// Mixed fleets: candidates disagree on their native path (e.g. DMR uses +// /anthropic/v1/messages while vLLM uses /v1/messages). Restricting the proxy +// to the full list would cause a 404 on whichever backend is selected but whose +// path was not used. Instead we return the largest path-compatible subset. +// Tie-break: prefer the subset whose path equals defaultPath; if no subset +// matches, use the path whose first endpoint appears earliest in the original +// endpoints slice. This ensures deterministic selection regardless of map +// iteration order. +// +// If profileLookup is nil or the endpoint list is empty, returns defaultPath and +// nil so the caller leaves everything unchanged. +func (a *Application) resolvePassthroughTargetPath(endpoints []*domain.Endpoint, defaultPath string) (string, []*domain.Endpoint) { if a.profileLookup == nil || len(endpoints) == 0 { - return defaultPath + return defaultPath, nil } - var resolvedPath string - for _, endpoint := range endpoints { - support := a.profileLookup.GetAnthropicSupport(endpoint.Type) - if support == nil || support.MessagesPath == "" { - return defaultPath + // Build path→endpoints groups while recording the first-seen index for each + // path. Iterating the original slice preserves endpoint order so the + // first-seen index is stable across calls. + pathGroups := make(map[string][]*domain.Endpoint, len(endpoints)) + firstSeen := make(map[string]int, len(endpoints)) // path → index of first endpoint in slice + for i, ep := range endpoints { + support := a.profileLookup.GetAnthropicSupport(ep.Type) + path := defaultPath + if support != nil && support.MessagesPath != "" { + path = support.MessagesPath } - if resolvedPath == "" { - resolvedPath = support.MessagesPath - continue + if _, exists := firstSeen[path]; !exists { + firstSeen[path] = i } - if support.MessagesPath != resolvedPath { - return defaultPath + pathGroups[path] = append(pathGroups[path], ep) + } + + // Uniform fleet: all candidates agree on a single path. + if len(pathGroups) == 1 { + for path := range pathGroups { + // nil filtered list signals "keep original slice" to the caller. + return path, nil } } - return resolvedPath + // Mixed fleet: pick the largest subset. Tie-break order: + // 1. defaultPath wins (avoids surprising path changes for the majority case). + // 2. The path whose first endpoint appears earliest in the original slice + // wins (deterministic, not dependent on map iteration order). + var bestPath string + var bestCount int + bestFirstSeen := -1 + + for path, group := range pathGroups { + count := len(group) + fs := firstSeen[path] + switch { + case count > bestCount: + bestPath, bestCount, bestFirstSeen = path, count, fs + case count == bestCount && path == defaultPath: + // defaultPath takes priority in a tie regardless of first-seen. + bestPath, bestFirstSeen = path, fs + case count == bestCount && bestPath != defaultPath && fs < bestFirstSeen: + // Neither candidate is defaultPath; earliest first-seen endpoint wins. + bestPath, bestFirstSeen = path, fs + } + } + + return bestPath, pathGroups[bestPath] } // executeTranslationRequest handles the translation path where requests are converted @@ -177,6 +227,19 @@ func (a *Application) executeTranslationRequest( a.logRequestStart(pr, len(endpoints)) + // For streaming requests, seed the context with an estimated input token count + // computed from the original (pre-translation) request body. This lets the + // streaming translator populate a non-zero input_tokens in message_start, matching + // the behaviour of vLLM and lmdeploy which both emit real input_tokens at that point. + // The estimate is overwritten by the actual upstream usage when it arrives. + if transformedReq.IsStreaming { + if estimator, ok := trans.(translator.InputTokenEstimator); ok && len(transformedReq.OriginalBody) > 0 { + estimate := estimator.EstimateInputTokens(transformedReq.OriginalBody) + ctx = context.WithValue(ctx, constants.ContextInputTokensKey, estimate) + r = r.WithContext(ctx) + } + } + // Execute proxy with appropriate response handling (streaming vs non-streaming) var proxyErr error if transformedReq.IsStreaming { @@ -215,13 +278,28 @@ func (a *Application) tryPassthrough( // Only pass endpoints whose backend natively supports the wire format. // Mixed deployments (e.g. ollama + vllm) must not block passthrough for - // the capable subset — the proxy will route within that filtered list. + // the capable subset - the proxy will route within that filtered list. + // + // Additionally, exclude endpoints that declare a limitation matching a + // feature present in this specific request. Token-counting limitations + // (e.g. token_counting_404) use different names and are never affected. + reqFeatures := detectRequestFeatures(bodyBytes) passthroughEndpoints := make([]*domain.Endpoint, 0, len(endpoints)) for _, ep := range endpoints { support := a.profileLookup.GetAnthropicSupport(ep.Type) - if support != nil && support.Enabled { - passthroughEndpoints = append(passthroughEndpoints, ep) + if support == nil || !support.Enabled { + continue + } + if reqFeatures.toolUse && support.HasLimitation(domain.AnthropicLimitationNoToolUse) { + continue + } + if reqFeatures.extendedThinking && support.HasLimitation(domain.AnthropicLimitationNoExtendedThinking) { + continue + } + if reqFeatures.vision && support.HasLimitation(domain.AnthropicLimitationNoVision) { + continue } + passthroughEndpoints = append(passthroughEndpoints, ep) } if !passthroughTrans.CanPassthrough(passthroughEndpoints, a.profileLookup) { @@ -331,7 +409,7 @@ func (a *Application) translationHandler(trans translator.RequestTranslator) htt return } - // Passthrough was not used — fall back to full translation. + // Passthrough was not used - fall back to full translation. mode := constants.TranslatorModeTranslation fallbackReason := a.resolveTranslationFallback(trans) @@ -369,17 +447,25 @@ func (a *Application) executeTranslatedNonStreamingRequest( return fmt.Errorf("proxy request failed: %w", err) } - // Parse OpenAI response + // Check for backend errors before attempting JSON parse. A reverse proxy or + // rate-limiter in front of the backend may return plain-text or HTML on 4xx/5xx, + // so parsing first would surface a misleading "failed to parse OpenAI response" + // error and lose the upstream status code (e.g. a plain-text 429 must stay 429). + if recorder.status >= 400 { + // Opportunistic parse: pass the map if the body is valid JSON so the + // error formatter can extract a message; pass nil otherwise and let + // extractAndLogBackendError fall back to its generic message. + var openaiErrResp map[string]interface{} + _ = json.Unmarshal(recorder.body.Bytes(), &openaiErrResp) + return a.handleNonStreamingBackendError(w, r, recorder, openaiErrResp, pr, trans) + } + + // Success path: strict parse - a malformed 200 body is a gateway error. var openaiResp map[string]interface{} if jerr := json.Unmarshal(recorder.body.Bytes(), &openaiResp); jerr != nil { return fmt.Errorf("failed to parse OpenAI response: %w", jerr) } - // handle backend errors - if recorder.status >= 400 { - return a.handleNonStreamingBackendError(w, r, recorder, openaiResp, pr, trans) - } - // transform and write successful response return a.writeTranslatedSuccessResponse(w, ctx, r, recorder, openaiResp, trans) } @@ -549,11 +635,20 @@ func (a *Application) executeTranslatedStreamingRequest( pipeReader, pipeWriter := io.Pipe() streamRecorder := newStreamingResponseRecorder(pipeWriter) + // Wrap the real client ResponseWriter so we can track whether any byte has + // actually been committed to the client (status line written or body written). + // This is distinct from streamRecorder.started, which only signals that the + // backend wrote into the pipe -- bytes may still be in the translator's buffer + // at that point, not yet on the wire to the client. + cw := newCommittedResponseWriter(w) + // run proxy in background while translation processes proxyErrChan := a.startProxyGoroutine(ctx, r, endpoints, pr, streamRecorder, pipeWriter) - // panic recovery prevents goroutine leak, cleanup before re-panic - defer a.handleStreamingPanic(pipeReader, pipeWriter, proxyErrChan, pr, trans) + // panic recovery prevents goroutine leak; sends a 502 (or SSE error event + // if the stream has already started) to the client rather than re-panicking, + // which would produce a connection reset with no useful error for the client. + defer a.handleStreamingPanic(cw, pipeReader, pipeWriter, proxyErrChan, cw, pr, trans) // Wait for headers before inspecting status. The select also handles context // cancellation so we don't block forever if the proxy errors without writing. @@ -566,18 +661,18 @@ func (a *Application) executeTranslatedStreamingRequest( // handle backend errors before starting sse stream if streamRecorder.status >= 400 { - a.handleStreamingBackendError(w, r, pipeReader, streamRecorder, proxyErrChan, pr, trans) + a.handleStreamingBackendError(cw, r, pipeReader, streamRecorder, proxyErrChan, pr, trans) return nil } // copy olla headers before stream starts - a.copyOllaHeaders(streamRecorder, w) - a.setModelHeaderIfMissing(w, pr.model) - // Write sticky headers before the first write to w commits the response. - a.setStickyResponseHeadersFromRequest(w, r) + a.copyOllaHeaders(streamRecorder, cw) + a.setModelHeaderIfMissing(cw, pr.model) + // Write sticky headers before the first write to cw commits the response. + a.setStickyResponseHeadersFromRequest(cw, r) // transform stream (blocks until done) and wait for proxy - return a.transformStreamAndWaitForProxy(ctx, pipeReader, w, r, proxyErrChan, trans) + return a.transformStreamAndWaitForProxy(ctx, pipeReader, cw, r, proxyErrChan, trans) } // writeStreamingNoEndpointsError writes error when no endpoints are available for streaming @@ -627,20 +722,31 @@ func (a *Application) startProxyGoroutine( return proxyErrChan } -// handleStreamingPanic recovers from panic during streaming to prevent goroutine leak +// handleStreamingPanic recovers from panic during streaming to prevent goroutine +// leak and connection resets. Instead of re-panicking (which terminates the +// connection abruptly and gives the client no useful error), it closes the pipe, +// drains the error channel, and writes an appropriate error to the client. +// +// The committed flag on cw tracks whether any byte has actually been written to +// the real client ResponseWriter (not just to the backend pipe). When committed, +// net/http silently ignores a WriteHeader(502), so we emit a spec-valid Anthropic +// SSE error event instead. When not committed, a plain HTTP 502 with a structured +// body reaches the client correctly. func (a *Application) handleStreamingPanic( + w http.ResponseWriter, pipeReader *io.PipeReader, pipeWriter *io.PipeWriter, proxyErrChan chan error, + cw *committedResponseWriter, pr *proxyRequest, trans translator.RequestTranslator, ) { if r := recover(); r != nil { - // Close both ends of the pipe to unblock the goroutine + // Close both ends of the pipe to unblock the proxy goroutine. pipeReader.Close() pipeWriter.Close() - // Drain the error channel to prevent goroutine leak + // Drain the buffered error channel to release the proxy goroutine. <-proxyErrChan a.logger.Error("Panic during stream transformation", @@ -648,8 +754,43 @@ func (a *Application) handleStreamingPanic( "translator", trans.Name(), "model", pr.model) - // Re-panic after cleanup to preserve the panic behavior - panic(r) + // If anything was already written to the real client, the response is + // committed and WriteHeader(502) is a no-op. Emitting a well-formed SSE + // error event is the only way to signal the failure without corrupting + // the stream with raw JSON. + if cw != nil && cw.committed.Load() { + a.writeSSEErrorEvent(w, "internal error during stream transformation") + return + } + + // Nothing reached the client yet - response is uncommitted, so a plain + // HTTP 502 with a structured error body reaches the client correctly. + if ew, ok := trans.(translator.ErrorWriter); ok { + ew.WriteError(w, errors.New("internal error during stream transformation"), http.StatusBadGateway) + } else { + http.Error(w, "internal error during stream transformation", http.StatusBadGateway) + } + } +} + +// writeSSEErrorEvent emits a spec-valid Anthropic SSE error event into an +// already-committed text/event-stream response. This is the only safe way to +// report a mid-stream failure; a WriteHeader call would be silently ignored by +// net/http once the response is committed. +func (a *Application) writeSSEErrorEvent(w http.ResponseWriter, message string) { + if w == nil { + return + } + payload := fmt.Sprintf(`{"type":"error","error":{"type":"api_error","message":%q}}`, message) + // Prefix with \n\n to guarantee a clean event boundary in case the panic + // occurred mid-write of a previous event. SSE parsers treat blank lines as + // field separators, so extra blank lines are harmless when already at a boundary. + sseEvent := fmt.Sprintf("\n\nevent: error\ndata: %s\n\n", payload) + _, _ = fmt.Fprint(w, sseEvent) + // Flush if the underlying writer supports it, so the client receives + // the event immediately rather than waiting for a buffer drain. + if f, ok := w.(http.Flusher); ok { + f.Flush() } } @@ -818,12 +959,32 @@ func (a *Application) tokenCountHandler(trans translator.RequestTranslator) http return } - // Write successful response + // Serialise before touching the response writer. WriteHeader(200) is a + // one-way door - once sent the client sees a 200 even if serialisation + // subsequently fails, which results in a truncated or empty body with a + // misleading success status. w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) - w.WriteHeader(http.StatusOK) - if err := json.NewEncoder(w).Encode(resp); err != nil { - a.logger.Error("Failed to encode token count response", "error", err) + if serialiser, ok := trans.(translator.TokenCountSerializer); ok { + body, serErr := serialiser.SerialiseCountTokens(resp) + if serErr != nil { + a.logger.Error("Failed to serialise token count response", "error", serErr) + if errorWriter, ok := trans.(translator.ErrorWriter); ok { + errorWriter.WriteError(w, serErr, http.StatusInternalServerError) + } else { + http.Error(w, "internal error serialising token count", http.StatusInternalServerError) + } + return + } + w.WriteHeader(http.StatusOK) + if _, wErr := w.Write(body); wErr != nil { //nolint:gosec // body is serialised JSON, not user-controlled data + a.logger.Error("Failed to write token count response", "error", wErr) + } + } else { + w.WriteHeader(http.StatusOK) + if err := json.NewEncoder(w).Encode(resp); err != nil { + a.logger.Error("Failed to encode token count response", "error", err) + } } } } @@ -882,6 +1043,100 @@ func (a *Application) recordTranslatorMetrics( a.statsCollector.RecordTranslatorRequest(event) } +// requestFeatures holds which Anthropic features are active in a specific request. +// Used to decide which capable endpoints are eligible for passthrough. +type requestFeatures struct { + toolUse bool + extendedThinking bool + vision bool +} + +// detectRequestFeatures inspects the raw Anthropic request body with gjson to +// determine which feature flags are active. It is deliberately cheap: gjson +// does not allocate a full parse tree, and we only scan the fields we care about. +// +// Content blocks can be a plain string (scalar) or an array of typed objects. +// The vision check handles both: a scalar string never contains image blocks, +// so we only iterate when content is a JSON array. +func detectRequestFeatures(body []byte) requestFeatures { + if len(body) == 0 { + return requestFeatures{} + } + + var f requestFeatures + + // Non-empty tools array → tool use is active. + if tools := gjson.GetBytes(body, "tools"); tools.IsArray() && len(tools.Array()) > 0 { + f.toolUse = true + } + + // Presence of top-level "thinking" key → extended thinking is active. + if gjson.GetBytes(body, "thinking").Exists() { + f.extendedThinking = true + } + + // Vision: any message content block with type=="image". + // messages.#.content yields the content field of every message. + // Each content entry may be a plain string (skip) or an array of blocks. + gjson.GetBytes(body, "messages.#.content").ForEach(func(_, contentVal gjson.Result) bool { + if !contentVal.IsArray() { + // Plain string content - no image blocks possible. + return true + } + contentVal.ForEach(func(_, block gjson.Result) bool { + if block.Get("type").Str == "image" { + f.vision = true + return false // stop iterating blocks + } + return true + }) + return !f.vision // stop iterating messages once found + }) + + return f +} + +// committedResponseWriter wraps the real client http.ResponseWriter and sets a +// committed flag on the first call to Write or WriteHeader. The panic handler +// reads this flag to decide whether to attempt a plain HTTP 502 (not yet committed) +// or inject an SSE error event into the already-open stream (committed). Using the +// real client write as the signal is correct; streamRecorder.started only tracks +// whether the backend wrote into the pipe, which can be true before any byte has +// reached the actual client. +type committedResponseWriter struct { + http.ResponseWriter + committed atomic.Bool +} + +func newCommittedResponseWriter(w http.ResponseWriter) *committedResponseWriter { + return &committedResponseWriter{ResponseWriter: w} +} + +func (c *committedResponseWriter) WriteHeader(statusCode int) { + c.committed.Store(true) + c.ResponseWriter.WriteHeader(statusCode) +} + +func (c *committedResponseWriter) Write(b []byte) (int, error) { + c.committed.Store(true) + return c.ResponseWriter.Write(b) +} + +// Flush marks the response as committed (bytes are en route to the client) and +// forwards the flush to the underlying writer. Without this, http.ResponseController +// cannot find http.Flusher on the wrapped writer, causing every SSE flush to return +// "feature not supported" and the translated streaming path to 502. +func (c *committedResponseWriter) Flush() { + c.committed.Store(true) + http.NewResponseController(c.ResponseWriter).Flush() //nolint:errcheck // best-effort flush; caller drives SSE cadence +} + +// Unwrap returns the underlying ResponseWriter so that http.ResponseController can +// reach optional interfaces (SetWriteDeadline, etc.) that we don't explicitly proxy. +func (c *committedResponseWriter) Unwrap() http.ResponseWriter { + return c.ResponseWriter +} + // abstract header access for both response types type headerGetter interface { Header() http.Header @@ -919,8 +1174,15 @@ type streamingResponseRecorder struct { writer io.Writer headers http.Header headersReady chan struct{} - closeOnce sync.Once status int + closeOnce sync.Once + // started is set true on the first Write, signalling that the upstream has + // begun emitting body bytes. The panic handler uses this to decide whether + // to inject an SSE error event into the already-open stream rather than + // attempting a plain HTTP 502 (which net/http silently ignores post-commit). + // atomic.Bool because Write is called from the proxy goroutine while + // handleStreamingPanic reads it from the handler goroutine. + started atomic.Bool } func newStreamingResponseRecorder(w io.Writer) *streamingResponseRecorder { @@ -937,12 +1199,13 @@ func (r *streamingResponseRecorder) Header() http.Header { } // ensureHeadersReady closes headersReady exactly once. It is safe to call from -// multiple goroutines and is idempotent — subsequent calls are no-ops. +// multiple goroutines and is idempotent - subsequent calls are no-ops. func (r *streamingResponseRecorder) ensureHeadersReady() { r.closeOnce.Do(func() { close(r.headersReady) }) } func (r *streamingResponseRecorder) Write(data []byte) (int, error) { + r.started.Store(true) r.ensureHeadersReady() return r.writer.Write(data) } @@ -954,6 +1217,6 @@ func (r *streamingResponseRecorder) WriteHeader(statusCode int) { } // Flush implements http.Flusher. The underlying io.Pipe is unbuffered -// (writes block until read), so there is nothing to flush — this is +// (writes block until read), so there is nothing to flush - this is // intentionally a no-op to satisfy http.ResponseController in proxy engines. func (r *streamingResponseRecorder) Flush() {} diff --git a/internal/app/handlers/handler_translation_passthrough_test.go b/internal/app/handlers/handler_translation_passthrough_test.go index 84f3ffaf..c0192406 100644 --- a/internal/app/handlers/handler_translation_passthrough_test.go +++ b/internal/app/handlers/handler_translation_passthrough_test.go @@ -1915,15 +1915,18 @@ func TestPassthrough_ProfileMessagesPathOverride(t *testing.T) { "proxy must use the backend's configured messages_path, not the translator default") } -// TestPassthrough_MixedMessagesPathsKeepsNeutralTargetPath verifies that the -// handler does not apply the first endpoint's backend-specific messages_path to -// every passthrough candidate. The proxy selector runs later, so mixed native -// Anthropic backends must keep the translator's neutral target path unless all -// candidates agree. -func TestPassthrough_MixedMessagesPathsKeepsNeutralTargetPath(t *testing.T) { +// TestPassthrough_MixedMessagesPathsFiltersToCompatibleSubset verifies that when +// capable backends disagree on their native Anthropic path, the handler restricts +// the proxy endpoint list to the path-compatible subset rather than passing all +// endpoints with a neutral path (which would cause 404s on mismatched backends). +// +// Fleet: DMR on /anthropic/v1/messages (1 endpoint) + vLLM on /v1/messages (1 endpoint). +// The translator's default path is /v1/messages, so on a size tie the /v1/messages +// group wins. The proxy receives only the vllm endpoint with path /v1/messages. +func TestPassthrough_MixedMessagesPathsFiltersToCompatibleSubset(t *testing.T) { t.Parallel() - var capturedEndpoint string + var capturedEndpoints []*domain.Endpoint var capturedPath string endpoints := []*domain.Endpoint{ @@ -1956,13 +1959,11 @@ func TestPassthrough_MixedMessagesPathsKeepsNeutralTargetPath(t *testing.T) { proxyService := &mockProxyService{ proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, eps []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { - require.Len(t, eps, 2) - selected := eps[1] - capturedEndpoint = selected.Name + capturedEndpoints = eps capturedPath = r.URL.Path w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) - w.Header().Set(constants.HeaderXOllaEndpoint, selected.Name) + w.Header().Set(constants.HeaderXOllaEndpoint, eps[0].Name) w.WriteHeader(http.StatusOK) return json.NewEncoder(w).Encode(map[string]interface{}{"type": "message", "id": "msg_vllm"}) }, @@ -1996,7 +1997,437 @@ func TestPassthrough_MixedMessagesPathsKeepsNeutralTargetPath(t *testing.T) { require.Equal(t, http.StatusOK, rec.Code) assert.Equal(t, string(constants.TranslatorModePassthrough), rec.Header().Get(constants.HeaderXOllaMode)) - assert.Equal(t, "vllm-backend", capturedEndpoint) + // On a tie the resolver prefers the subset whose path matches the translator + // default (/v1/messages), which is the vllm group. + require.Len(t, capturedEndpoints, 1, + "proxy must receive only the path-compatible subset, not the full mixed fleet") + assert.Equal(t, "vllm-backend", capturedEndpoints[0].Name, + "vllm endpoint (path=/v1/messages, matches translator default) must be selected") assert.Equal(t, "/v1/messages", capturedPath, - "proxy must not use DMR's first-endpoint path when the selected backend uses /v1/messages") + "request path must match the selected subset's native messages path") +} + +// TestPassthrough_MixedFleet_TwoVllmOneDmr_PicksVllmSubset verifies that when the +// majority of capable backends share a path, that larger subset wins regardless of +// the tie-break rule, and the DMR endpoint is excluded. +func TestPassthrough_MixedFleet_TwoVllmOneDmr_PicksVllmSubset(t *testing.T) { + t.Parallel() + + var capturedEndpoints []*domain.Endpoint + var capturedPath string + + endpoints := []*domain.Endpoint{ + {Name: "vllm-1", Type: "vllm", Status: domain.StatusHealthy, URLString: "http://localhost:8001"}, + {Name: "vllm-2", Type: "vllm", Status: domain.StatusHealthy, URLString: "http://localhost:8002"}, + {Name: "dmr-1", Type: "docker-model-runner", Status: domain.StatusHealthy, URLString: "http://localhost:12434"}, + } + + profileLookup := &mockPassthroughProfileLookup{ + configs: map[string]*domain.AnthropicSupportConfig{ + "vllm": {Enabled: true, MessagesPath: "/v1/messages"}, + "docker-model-runner": {Enabled: true, MessagesPath: "/anthropic/v1/messages"}, + }, + } + + trans := &mockPassthroughTranslator{ + name: "anthropic", + passthroughEnabled: true, + profileLookup: profileLookup, + } + + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, eps []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + capturedEndpoints = eps + capturedPath = r.URL.Path + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(http.StatusOK) + return json.NewEncoder(w).Encode(map[string]interface{}{"type": "message", "id": "msg_1"}) + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{getEndpointsFunc: func() []*domain.Endpoint { return endpoints }}, + inspectorChain: inspector.NewChain(&mockStyledLogger{}), + profileFactory: &mockProfileFactory{}, + profileLookup: profileLookup, + discoveryService: &mockDiscoveryServiceWithEndpoints{endpoints: endpoints}, + Config: &config.Config{}, + } + + handler := app.translationHandler(trans) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "claude-3-5-sonnet-20241022", + "messages": []map[string]interface{}{{"role": "user", "content": "hello"}}, + }) + + req := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqBody)) + req.Header.Set(constants.HeaderContentType, constants.ContentTypeJSON) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, string(constants.TranslatorModePassthrough), rec.Header().Get(constants.HeaderXOllaMode)) + // The vllm group has 2 endpoints vs DMR's 1 — larger subset wins. + require.Len(t, capturedEndpoints, 2, + "both vllm endpoints must be passed; DMR must be excluded") + for _, ep := range capturedEndpoints { + assert.Equal(t, "vllm", ep.Type, "all passed endpoints must be vllm") + } + assert.Equal(t, "/v1/messages", capturedPath, + "path must be the vllm native path, not the DMR path") +} + +// TestPassthrough_UniformDmrFleet_AllEndpointsPassed verifies that a uniform +// DMR-only fleet (all on /anthropic/v1/messages) passes all endpoints through +// without filtering and uses the DMR-specific path. +func TestPassthrough_UniformDmrFleet_AllEndpointsPassed(t *testing.T) { + t.Parallel() + + var capturedEndpoints []*domain.Endpoint + var capturedPath string + + endpoints := []*domain.Endpoint{ + {Name: "dmr-1", Type: "docker-model-runner", Status: domain.StatusHealthy, URLString: "http://localhost:12434"}, + {Name: "dmr-2", Type: "docker-model-runner", Status: domain.StatusHealthy, URLString: "http://localhost:12435"}, + } + + profileLookup := &mockPassthroughProfileLookup{ + configs: map[string]*domain.AnthropicSupportConfig{ + "docker-model-runner": {Enabled: true, MessagesPath: "/anthropic/v1/messages"}, + }, + } + + trans := &mockPassthroughTranslator{ + name: "anthropic", + passthroughEnabled: true, + profileLookup: profileLookup, + } + + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, eps []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + capturedEndpoints = eps + capturedPath = r.URL.Path + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(http.StatusOK) + return json.NewEncoder(w).Encode(map[string]interface{}{"type": "message", "id": "msg_dmr"}) + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{getEndpointsFunc: func() []*domain.Endpoint { return endpoints }}, + inspectorChain: inspector.NewChain(&mockStyledLogger{}), + profileFactory: &mockProfileFactory{}, + profileLookup: profileLookup, + discoveryService: &mockDiscoveryServiceWithEndpoints{endpoints: endpoints}, + Config: &config.Config{}, + } + + handler := app.translationHandler(trans) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "ai/llama3.2", + "messages": []map[string]interface{}{{"role": "user", "content": "hello"}}, + }) + + req := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqBody)) + req.Header.Set(constants.HeaderContentType, constants.ContentTypeJSON) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, string(constants.TranslatorModePassthrough), rec.Header().Get(constants.HeaderXOllaMode)) + // Uniform fleet: no filtering — both DMR endpoints must be passed. + require.Len(t, capturedEndpoints, 2, + "uniform DMR fleet must pass all endpoints without filtering") + assert.Equal(t, "/anthropic/v1/messages", capturedPath, + "path must be the DMR native path for a uniform DMR fleet") +} + +// TestPassthrough_LimitationFiltering_NoToolUse verifies that an endpoint declaring +// no_tool_use is excluded from the passthrough subset when the request contains tools, +// but included when the request has no tools. +func TestPassthrough_LimitationFiltering_NoToolUse(t *testing.T) { + t.Parallel() + + endpoints := []*domain.Endpoint{ + {Name: "limited-backend", Type: "limited", Status: domain.StatusHealthy}, + } + + profileLookup := &mockPassthroughProfileLookup{ + configs: map[string]*domain.AnthropicSupportConfig{ + "limited": { + Enabled: true, + MessagesPath: "/v1/messages", + Limitations: []string{domain.AnthropicLimitationNoToolUse}, + }, + }, + } + + trans := &mockPassthroughTranslator{ + name: "anthropic", + passthroughEnabled: true, + profileLookup: profileLookup, + transformRequestFunc: func(ctx context.Context, r *http.Request) (*translator.TransformedRequest, error) { + return &translator.TransformedRequest{ + OpenAIRequest: map[string]interface{}{"model": "claude-3-5-sonnet-20241022", "messages": []interface{}{}}, + ModelName: "claude-3-5-sonnet-20241022", + IsStreaming: false, + TargetPath: "/v1/chat/completions", + }, nil + }, + transformResponseFunc: func(ctx context.Context, openaiResp interface{}, original *http.Request) (interface{}, error) { + return map[string]interface{}{"type": "message", "id": "msg_translated"}, nil + }, + implementsErrorWriter: true, + } + + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, eps []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(http.StatusOK) + return json.NewEncoder(w).Encode(map[string]interface{}{"type": "message", "id": "msg_proxy"}) + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{getEndpointsFunc: func() []*domain.Endpoint { return endpoints }}, + inspectorChain: inspector.NewChain(&mockStyledLogger{}), + profileFactory: &mockProfileFactory{}, + profileLookup: profileLookup, + discoveryService: &mockDiscoveryServiceWithEndpoints{endpoints: endpoints}, + Config: &config.Config{}, + } + + handler := app.translationHandler(trans) + + // Request WITH tools — endpoint must be excluded, falls back to translation. + reqWithTools, _ := json.Marshal(map[string]interface{}{ + "model": "claude-3-5-sonnet-20241022", + "messages": []map[string]interface{}{{"role": "user", "content": "use a tool"}}, + "tools": []map[string]interface{}{ + {"name": "get_weather", "description": "Get weather", "input_schema": map[string]interface{}{"type": "object"}}, + }, + }) + + req := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqWithTools)) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusOK, rec.Code) + assert.Empty(t, rec.Header().Get("X-Olla-Mode"), + "endpoint with no_tool_use must be excluded when request uses tools — falls back to translation") + + // Request WITHOUT tools — endpoint must be included, passthrough used. + reqWithoutTools, _ := json.Marshal(map[string]interface{}{ + "model": "claude-3-5-sonnet-20241022", + "messages": []map[string]interface{}{{"role": "user", "content": "hello"}}, + }) + + req2 := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqWithoutTools)) + rec2 := httptest.NewRecorder() + handler.ServeHTTP(rec2, req2) + + assert.Equal(t, http.StatusOK, rec2.Code) + assert.Equal(t, "passthrough", rec2.Header().Get("X-Olla-Mode"), + "endpoint with no_tool_use must be included when request has no tools") +} + +// TestPassthrough_LimitationFiltering_TokenCounting verifies that token-counting +// limitations (token_counting_404, no_token_counting) never affect messages passthrough. +func TestPassthrough_LimitationFiltering_TokenCounting(t *testing.T) { + t.Parallel() + + var proxyCalled bool + + endpoints := []*domain.Endpoint{ + {Name: "tc-limited", Type: "tc_backend", Status: domain.StatusHealthy}, + } + + profileLookup := &mockPassthroughProfileLookup{ + configs: map[string]*domain.AnthropicSupportConfig{ + "tc_backend": { + Enabled: true, + MessagesPath: "/v1/messages", + // Token-counting limitations must never affect messages passthrough. + Limitations: []string{"token_counting_404", "no_token_counting"}, + }, + }, + } + + trans := &mockPassthroughTranslator{ + name: "anthropic", + passthroughEnabled: true, + profileLookup: profileLookup, + } + + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, eps []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + proxyCalled = true + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(http.StatusOK) + return json.NewEncoder(w).Encode(map[string]interface{}{"type": "message", "id": "msg_tc"}) + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{getEndpointsFunc: func() []*domain.Endpoint { return endpoints }}, + inspectorChain: inspector.NewChain(&mockStyledLogger{}), + profileFactory: &mockProfileFactory{}, + profileLookup: profileLookup, + discoveryService: &mockDiscoveryServiceWithEndpoints{endpoints: endpoints}, + Config: &config.Config{}, + } + + handler := app.translationHandler(trans) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "claude-3-5-sonnet-20241022", + "messages": []map[string]interface{}{{"role": "user", "content": "hello"}}, + }) + + req := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqBody)) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusOK, rec.Code) + assert.True(t, proxyCalled, "proxy must be called — token-counting limitations must not exclude the endpoint") + assert.Equal(t, "passthrough", rec.Header().Get("X-Olla-Mode"), + "token-counting limitations must never block messages passthrough") +} + +// TestPassthrough_LimitationFiltering_MixedFleet verifies that in a mixed fleet one +// endpoint is excluded by a feature limitation while the other passes through. +// vllm declares no_tool_use; llamacpp has no limitations. A tool-use request must +// route only to llamacpp via passthrough. +func TestPassthrough_LimitationFiltering_MixedFleet(t *testing.T) { + t.Parallel() + + var capturedEndpoints []*domain.Endpoint + + endpoints := []*domain.Endpoint{ + {Name: "vllm-1", Type: "vllm", Status: domain.StatusHealthy}, + {Name: "llamacpp-1", Type: "llamacpp", Status: domain.StatusHealthy}, + } + + profileLookup := &mockPassthroughProfileLookup{ + configs: map[string]*domain.AnthropicSupportConfig{ + "vllm": { + Enabled: true, + MessagesPath: "/v1/messages", + Limitations: []string{domain.AnthropicLimitationNoToolUse}, + }, + "llamacpp": { + Enabled: true, + MessagesPath: "/v1/messages", + // No limitations — supports tool use. + }, + }, + } + + trans := &mockPassthroughTranslator{ + name: "anthropic", + passthroughEnabled: true, + profileLookup: profileLookup, + } + + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, eps []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + capturedEndpoints = eps + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(http.StatusOK) + return json.NewEncoder(w).Encode(map[string]interface{}{"type": "message", "id": "msg_mixed"}) + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{getEndpointsFunc: func() []*domain.Endpoint { return endpoints }}, + inspectorChain: inspector.NewChain(&mockStyledLogger{}), + profileFactory: &mockProfileFactory{}, + profileLookup: profileLookup, + discoveryService: &mockDiscoveryServiceWithEndpoints{endpoints: endpoints}, + Config: &config.Config{}, + } + + handler := app.translationHandler(trans) + + reqWithTools, _ := json.Marshal(map[string]interface{}{ + "model": "claude-3-5-sonnet-20241022", + "messages": []map[string]interface{}{{"role": "user", "content": "use a tool"}}, + "tools": []map[string]interface{}{ + {"name": "search", "description": "Search", "input_schema": map[string]interface{}{"type": "object"}}, + }, + }) + + req := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqWithTools)) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "passthrough", rec.Header().Get("X-Olla-Mode"), + "llamacpp (no limitations) must accept the tool-use request via passthrough") + require.Len(t, capturedEndpoints, 1, "only llamacpp must be in the passthrough subset") + assert.Equal(t, "llamacpp-1", capturedEndpoints[0].Name, + "vllm (no_tool_use) must be excluded; llamacpp must handle the request") +} + +// TestResolvePassthroughTargetPath_TieBreakDeterminism ensures that when two equal-sized +// non-default path groups exist, the function always resolves to the path whose first +// endpoint appears earliest in the input slice — never the other group, regardless of +// Go's randomised map iteration order. +func TestResolvePassthroughTargetPath_TieBreakDeterminism(t *testing.T) { + t.Parallel() + + // Two backend types with distinct native Anthropic paths, neither matching defaultPath. + // Each group has exactly one endpoint so the sizes are equal — a pure tie. + // The endpoint for "alpha" appears before "beta" in the slice, so "alpha" must always win. + const defaultPath = "/v1/messages" + const alphaPath = "/alpha/v1/messages" + const betaPath = "/beta/v1/messages" + + profileLookup := &mockPassthroughProfileLookup{ + configs: map[string]*domain.AnthropicSupportConfig{ + "alpha": {Enabled: true, MessagesPath: alphaPath}, + "beta": {Enabled: true, MessagesPath: betaPath}, + }, + } + + endpoints := []*domain.Endpoint{ + {Name: "alpha-1", Type: "alpha", Status: domain.StatusHealthy}, + {Name: "beta-1", Type: "beta", Status: domain.StatusHealthy}, + } + + app := &Application{ + profileLookup: profileLookup, + } + + // Run enough iterations to expose any map-order flapping. + const iterations = 20 + for i := range iterations { + resolvedPath, filteredEndpoints := app.resolvePassthroughTargetPath(endpoints, defaultPath) + if resolvedPath != alphaPath { + t.Errorf("iteration %d: expected %q (first-seen path), got %q", i, alphaPath, resolvedPath) + } + // The filtered set must contain only the alpha endpoint. + if len(filteredEndpoints) != 1 || filteredEndpoints[0].Name != "alpha-1" { + t.Errorf("iteration %d: expected [alpha-1], got %v", i, filteredEndpoints) + } + } } diff --git a/internal/app/handlers/handler_translation_test.go b/internal/app/handlers/handler_translation_test.go index 0f032d29..7dde0c25 100644 --- a/internal/app/handlers/handler_translation_test.go +++ b/internal/app/handlers/handler_translation_test.go @@ -27,14 +27,18 @@ import ( // mockTranslator implements RequestTranslator for testing type mockTranslator struct { - name string - transformRequestFunc func(ctx context.Context, r *http.Request) (*translator.TransformedRequest, error) - transformResponseFunc func(ctx context.Context, openaiResp interface{}, original *http.Request) (interface{}, error) - transformStreamingFunc func(ctx context.Context, openaiStream io.Reader, w http.ResponseWriter, original *http.Request) error - writeErrorFunc func(w http.ResponseWriter, err error, statusCode int) - pathProvider string - implementsErrorWriter bool - implementsPathProvider bool + name string + transformRequestFunc func(ctx context.Context, r *http.Request) (*translator.TransformedRequest, error) + transformResponseFunc func(ctx context.Context, openaiResp interface{}, original *http.Request) (interface{}, error) + transformStreamingFunc func(ctx context.Context, openaiStream io.Reader, w http.ResponseWriter, original *http.Request) error + writeErrorFunc func(w http.ResponseWriter, err error, statusCode int) + countTokensFunc func(ctx context.Context, r *http.Request) (*translator.TokenCountResponse, error) + serialiseCountTokensFunc func(resp *translator.TokenCountResponse) ([]byte, error) + pathProvider string + implementsErrorWriter bool + implementsPathProvider bool + implementsTokenCounter bool + implementsTokenSerializer bool } func (m *mockTranslator) Name() string { @@ -94,6 +98,20 @@ func (m *mockTranslator) WriteError(w http.ResponseWriter, err error, statusCode panic("WriteError called on translator that doesn't implement ErrorWriter") } +func (m *mockTranslator) CountTokens(ctx context.Context, r *http.Request) (*translator.TokenCountResponse, error) { + if m.implementsTokenCounter && m.countTokensFunc != nil { + return m.countTokensFunc(ctx, r) + } + panic("CountTokens called on translator that doesn't implement TokenCounter") +} + +func (m *mockTranslator) SerialiseCountTokens(resp *translator.TokenCountResponse) ([]byte, error) { + if m.implementsTokenSerializer && m.serialiseCountTokensFunc != nil { + return m.serialiseCountTokensFunc(resp) + } + panic("SerialiseCountTokens called on translator that doesn't implement TokenCountSerializer") +} + // mockProxyService implements ProxyService for testing type mockProxyService struct { proxyFunc func(ctx context.Context, w http.ResponseWriter, r *http.Request, endpoints []*domain.Endpoint, stats *ports.RequestStats, logger logger.StyledLogger) error @@ -136,6 +154,128 @@ func (m *mockProxyService) ProxyRequestToEndpoints( return json.NewEncoder(w).Encode(response) } +// TestTranslationHandler_PlainTextBackendError verifies that when a backend returns +// a non-JSON 4xx/5xx body (e.g. a plain-text 429 from a rate-limiter), the handler +// preserves the upstream status code and emits a structured Anthropic error rather +// than surfacing a misleading "failed to parse OpenAI response" 502. +func TestTranslationHandler_PlainTextBackendError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + backendStatus int + backendBody string + backendCT string + expectedStatus int + expectedType string + }{ + { + name: "plain_text_429", + backendStatus: http.StatusTooManyRequests, + backendBody: "Rate limit exceeded. Please try again later.", + backendCT: "text/plain", + expectedStatus: http.StatusTooManyRequests, + expectedType: "rate_limit_error", + }, + { + name: "html_503", + backendStatus: http.StatusServiceUnavailable, + backendBody: "Service Unavailable", + backendCT: "text/html", + expectedStatus: http.StatusServiceUnavailable, + expectedType: "overloaded_error", // mock writeErrorFunc maps 503 → overloaded_error + }, + { + name: "json_400_preserved", + backendStatus: http.StatusBadRequest, + backendBody: `{"error":{"message":"Invalid temperature value","type":"invalid_request_error"}}`, + backendCT: "application/json", + // Existing behaviour must be preserved: message extracted from JSON. + expectedStatus: http.StatusBadRequest, + expectedType: "invalid_request_error", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + mockLogger := &mockStyledLogger{} + + trans := &mockTranslator{ + name: "anthropic", + implementsErrorWriter: true, + writeErrorFunc: func(w http.ResponseWriter, err error, statusCode int) { + errorType := "api_error" + switch statusCode { + case http.StatusBadRequest: + errorType = "invalid_request_error" + case http.StatusTooManyRequests: + errorType = "rate_limit_error" + case http.StatusServiceUnavailable: + errorType = "overloaded_error" + } + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(statusCode) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "type": "error", + "error": map[string]interface{}{ + "type": errorType, + "message": err.Error(), + }, + }) + }, + } + + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, endpoints []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + w.Header().Set(constants.HeaderContentType, tt.backendCT) + w.WriteHeader(tt.backendStatus) + _, err := w.Write([]byte(tt.backendBody)) + return err + }, + } + + app := &Application{ + logger: mockLogger, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{}, + inspectorChain: inspector.NewChain(mockLogger), + profileFactory: &mockProfileFactory{}, + discoveryService: &mockDiscoveryServiceForTranslation{}, + Config: &config.Config{}, + } + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "test-model", + "max_tokens": 100, + "messages": []map[string]interface{}{{"role": "user", "content": "hello"}}, + }) + + req := httptest.NewRequest("POST", "/olla/anthropic/v1/messages", bytes.NewReader(reqBody)) + req.Header.Set(constants.HeaderContentType, constants.ContentTypeJSON) + rec := httptest.NewRecorder() + + handler := app.translationHandler(trans) + handler.ServeHTTP(rec, req) + + assert.Equal(t, tt.expectedStatus, rec.Code, + "upstream status code must be preserved, not flattened to 502") + + var errBody map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &errBody), + "error body must be valid JSON") + assert.Equal(t, "error", errBody["type"]) + + errObj, ok := errBody["error"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, tt.expectedType, errObj["type"], + "error type must match the expected Anthropic error type for this status") + }) + } +} + func TestTranslationHandler_NonStreaming(t *testing.T) { mockLogger := &mockStyledLogger{} trans := &mockTranslator{ @@ -704,6 +844,8 @@ func (m *mockTranslatorWithoutErrorWriter) TransformStreamingResponse(ctx contex } func TestTranslationHandler_StreamingPanicRecovery(t *testing.T) { + t.Parallel() + mockLogger := &mockStyledLogger{} panicTrans := &mockTranslator{ @@ -712,7 +854,7 @@ func TestTranslationHandler_StreamingPanicRecovery(t *testing.T) { writeErrorFunc: func(w http.ResponseWriter, err error, statusCode int) { w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) w.WriteHeader(statusCode) - json.NewEncoder(w).Encode(map[string]interface{}{ + _ = json.NewEncoder(w).Encode(map[string]interface{}{ "error": err.Error(), }) }, @@ -781,18 +923,24 @@ func TestTranslationHandler_StreamingPanicRecovery(t *testing.T) { req := httptest.NewRequest("POST", "/test", bytes.NewReader(body)) rec := httptest.NewRecorder() - defer func() { - r := recover() - if r == nil { - t.Fatal("Expected panic to be propagated after cleanup") - } - - panicMsg := fmt.Sprintf("%v", r) - assert.Contains(t, panicMsg, "simulated panic during stream transformation", - "Panic message should contain original panic text") - }() - + // The handler must NOT re-panic. A re-panic resets the TCP connection, + // giving the client no indication of what went wrong. + // + // The proxy writes one SSE event into the pipe, but the translator panics + // immediately without reading from it and writing anything to the real client + // ResponseWriter. The committedResponseWriter therefore has committed=false, + // and the handler correctly falls back to a plain HTTP 502 rather than an SSE + // error event. This is the correct behaviour: no bytes reached the client, so + // the response is not committed and a proper error status can be sent. handler.ServeHTTP(rec, req) + + // The panic must not propagate as a connection reset. Because the translator + // panicked before writing to the client, the response is uncommitted and a + // 502 with a structured error body is the correct outcome. + assert.Equal(t, http.StatusBadGateway, rec.Code, + "status must be 502 when panic fires before any bytes reach the client") + assert.NotContains(t, rec.Body.String(), "event: error", + "no SSE event should appear when the response was not committed") } // Verifies that: @@ -1570,7 +1718,7 @@ func TestExecuteTranslatedStreamingRequest_ProxyErrorBeforeWrite(t *testing.T) { select { case <-done: - // Handler returned — no deadlock. The response status must indicate an error. + // Handler returned - no deadlock. The response status must indicate an error. assert.GreaterOrEqual(t, rec.Code, http.StatusBadRequest, "expected an error status when proxy fails before writing headers") case <-ctx.Done(): @@ -1657,7 +1805,7 @@ func TestExecuteTranslatedStreamingRequest_ContextCancellationUnblocks(t *testin // correctly unblocks the select, the handler returns well before this deadline. select { case <-done: - // Returned after cancellation — correct behaviour. + // Returned after cancellation - correct behaviour. case <-time.After(3 * time.Second): t.Fatal("handler did not unblock after context cancellation") } @@ -1756,3 +1904,354 @@ func TestExecuteTranslatedStreamingRequest_SuccessfulFlow(t *testing.T) { assert.NotEmpty(t, rec.Header().Get(constants.HeaderXOllaRequestID), "X-Olla-Request-ID should be copied to the client response") } + +// TestHandleStreamingPanic_Writes502 verifies that a panic inside the streaming +// path produces a 502 response to the client rather than re-panicking (which +// would close the connection without sending any error). +func TestHandleStreamingPanic_Writes502(t *testing.T) { + t.Parallel() + + mockLogger := &mockStyledLogger{} + app := &Application{logger: mockLogger} + + trans := &mockTranslator{ + name: "panicking-translator", + implementsErrorWriter: true, + writeErrorFunc: func(w http.ResponseWriter, err error, statusCode int) { + w.WriteHeader(statusCode) + }, + } + + // Set up a pipe and a buffered error channel as the real code does. + pipeReader, pipeWriter := io.Pipe() + proxyErrChan := make(chan error, 1) + + // Pre-fill the error channel so the drain in handleStreamingPanic does not block. + proxyErrChan <- nil + + rec := httptest.NewRecorder() + + // committedResponseWriter with committed=false simulates a panic before any + // byte was written to the real client - response is uncommitted. + cw := newCommittedResponseWriter(rec) + + // Call handleStreamingPanic directly, simulating the deferred call after a panic. + // We wrap in a function that sets up a panic so recover() fires. + func() { + defer app.handleStreamingPanic(cw, pipeReader, pipeWriter, proxyErrChan, cw, &proxyRequest{ + requestLogger: mockLogger, + }, trans) + panic("simulated streaming panic") + }() + + // The panic must NOT propagate - if it did, the test would fail here. + assert.Equal(t, http.StatusBadGateway, rec.Code, + "streaming panic must produce 502 not a connection reset") +} + +// TestHandleStreamingPanic_AfterStreamStarted_EmitsSSEError verifies that when a +// panic fires after the upstream has already written body bytes, the handler emits +// a well-formed Anthropic SSE error event rather than attempting a plain HTTP 502 +// (which net/http silently ignores once the response is committed). +func TestHandleStreamingPanic_AfterStreamStarted_EmitsSSEError(t *testing.T) { + t.Parallel() + + mockLogger := &mockStyledLogger{} + app := &Application{logger: mockLogger} + + trans := &mockTranslator{ + name: "panicking-translator", + implementsErrorWriter: true, + writeErrorFunc: func(w http.ResponseWriter, err error, statusCode int) { + w.WriteHeader(statusCode) + }, + } + + pipeReader, pipeWriter := io.Pipe() + proxyErrChan := make(chan error, 1) + proxyErrChan <- nil + + rec := httptest.NewRecorder() + + // committedResponseWriter with committed=true simulates the state after at + // least one byte has been written to the real client ResponseWriter. + cw := newCommittedResponseWriter(rec) + cw.committed.Store(true) + + func() { + defer app.handleStreamingPanic(cw, pipeReader, pipeWriter, proxyErrChan, cw, &proxyRequest{ + requestLogger: mockLogger, + }, trans) + panic("simulated mid-stream panic") + }() + + body := rec.Body.String() + // Status must remain 200 - WriteHeader(502) is a no-op after the first Write. + assert.Equal(t, http.StatusOK, rec.Code, + "status must stay 200 when the panic fires after the stream has started") + assert.Contains(t, body, "event: error", + "a spec-valid SSE error line must be appended mid-stream") + assert.Contains(t, body, "data: ", + "error event must include a data line") + assert.Contains(t, body, "api_error", + "Anthropic error type must appear in the SSE data payload") + assert.Contains(t, body, "internal error during stream transformation", + "error message must appear in the SSE data payload") + // Ensure no raw JSON is injected outside the SSE envelope. + assert.NotContains(t, body, `{"error":`, + "raw JSON error must not be injected outside the SSE envelope") +} + +// TestStreamingPanic_BeforeWrite_Returns502 exercises the full translation handler +// end-to-end with a panic that fires before any SSE data is written to the client. +// The response must be HTTP 502 with an Anthropic-formatted JSON error body. +func TestStreamingPanic_BeforeWrite_Returns502(t *testing.T) { + t.Parallel() + + mockLogger := &mockStyledLogger{} + + // The translator panics inside TransformStreamingResponse before writing + // anything to w. Because no bytes reach the client, the response is + // uncommitted and a 502 can be written. + panicTrans := &mockTranslator{ + name: "early-panic-translator", + implementsErrorWriter: true, + writeErrorFunc: func(w http.ResponseWriter, err error, statusCode int) { + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(statusCode) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "type": "error", + "error": map[string]interface{}{ + "type": "api_error", + "message": err.Error(), + }, + }) + }, + transformRequestFunc: func(ctx context.Context, r *http.Request) (*translator.TransformedRequest, error) { + return &translator.TransformedRequest{ + OpenAIRequest: map[string]interface{}{ + "model": "test-model", + "stream": true, + "messages": []interface{}{ + map[string]interface{}{"role": "user", "content": "test"}, + }, + }, + ModelName: "test-model", + IsStreaming: true, + }, nil + }, + transformStreamingFunc: func(ctx context.Context, openaiStream io.Reader, w http.ResponseWriter, original *http.Request) error { + // Panic before writing anything to w - stream is not started. + panic("panic before any write") + }, + } + + // Proxy does NOT write any body data (no Write call, only WriteHeader), so + // streamRecorder.started remains false when the panic fires. + proxyService := &mockProxyService{ + proxyFunc: func(ctx context.Context, w http.ResponseWriter, r *http.Request, endpoints []*domain.Endpoint, stats *ports.RequestStats, rlog logger.StyledLogger) error { + w.Header().Set(constants.HeaderContentType, "text/event-stream") + w.WriteHeader(http.StatusOK) + // No Write call - started remains false. + return nil + }, + } + + app := &Application{ + logger: mockLogger, + proxyService: proxyService, + statsCollector: &mockStatsCollector{}, + repository: &mockEndpointRepository{}, + inspectorChain: inspector.NewChain(mockLogger), + profileFactory: &mockProfileFactory{}, + discoveryService: &mockDiscoveryServiceForTranslation{}, + Config: &config.Config{}, + } + + handler := app.translationHandler(panicTrans) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "test-model", + "stream": true, + "messages": []interface{}{ + map[string]interface{}{"role": "user", "content": "test"}, + }, + }) + req := httptest.NewRequest("POST", "/test", bytes.NewReader(reqBody)) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusBadGateway, rec.Code, + "502 must be returned when the panic fires before any bytes are written") + + var errBody map[string]interface{} + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &errBody), + "error body must be valid JSON") + assert.Equal(t, "error", errBody["type"], + "Anthropic error envelope type must be 'error'") + errObj, ok := errBody["error"].(map[string]interface{}) + require.True(t, ok, "error object must be present") + assert.Equal(t, "api_error", errObj["type"]) +} + +// TestCommittedResponseWriter_FlushSetsCommitted verifies that Flush() marks the response +// as committed. Without this the translated SSE path calls http.NewResponseController(cw).Flush() +// which can't find http.Flusher on the wrapper, returns "feature not supported", and the +// streaming handler treats the flush failure as a transform error (502). +func TestCommittedResponseWriter_FlushSetsCommitted(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + cw := newCommittedResponseWriter(rec) + + if cw.committed.Load() { + t.Fatal("committed must be false before any write") + } + + rc := http.NewResponseController(cw) + if err := rc.Flush(); err != nil { + t.Fatalf("Flush via ResponseController must succeed, got: %v", err) + } + + if !cw.committed.Load() { + t.Fatal("committed must be true after Flush()") + } +} + +// TestCommittedResponseWriter_Unwrap verifies that the underlying ResponseWriter is +// reachable via Unwrap so ResponseController can discover optional interfaces +// (SetWriteDeadline, etc.) beyond what committedResponseWriter explicitly proxies. +func TestCommittedResponseWriter_Unwrap(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + cw := newCommittedResponseWriter(rec) + + if cw.Unwrap() != rec { + t.Fatal("Unwrap must return the exact underlying ResponseWriter") + } +} + +// TestTokenCountHandler_SerialiseFailureReturns500 verifies that when a translator +// implements TokenCountSerializer but SerialiseCountTokens returns an error, the +// handler must NOT commit a 200 before the failure. Before the fix, WriteHeader(200) +// was called unconditionally before serialisation, so the client received a 200 with +// an empty or partial body. +func TestTokenCountHandler_SerialiseFailureReturns500(t *testing.T) { + t.Parallel() + + mockLogger := &mockStyledLogger{} + serErr := errors.New("json serialisation kaboom") + + trans := &mockTranslator{ + name: "anthropic", + implementsErrorWriter: true, + implementsTokenCounter: true, + implementsTokenSerializer: true, + writeErrorFunc: func(w http.ResponseWriter, err error, statusCode int) { + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(statusCode) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "type": "error", + "error": map[string]interface{}{ + "type": "api_error", + "message": err.Error(), + }, + }) + }, + countTokensFunc: func(_ context.Context, _ *http.Request) (*translator.TokenCountResponse, error) { + return &translator.TokenCountResponse{InputTokens: 42}, nil + }, + serialiseCountTokensFunc: func(_ *translator.TokenCountResponse) ([]byte, error) { + return nil, serErr + }, + } + + app := &Application{ + logger: mockLogger, + } + + handler := app.tokenCountHandler(trans) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "test-model", + "messages": []interface{}{map[string]interface{}{"role": "user", "content": "hi"}}, + }) + req := httptest.NewRequest(http.MethodPost, "/olla/anthropic/v1/messages/count_tokens", bytes.NewReader(reqBody)) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + // The status must not be 200 - serialisation failed before any bytes were committed. + if rec.Code == http.StatusOK { + t.Errorf("expected non-200 status when serialisation fails, got %d (body: %s)", rec.Code, rec.Body.String()) + } + if rec.Code != http.StatusInternalServerError { + t.Errorf("expected 500, got %d", rec.Code) + } +} + +// TestTokenCountHandler_SerialiseSuccess verifies the happy path: successful +// serialisation writes the body with a 200 status. +func TestTokenCountHandler_SerialiseSuccess(t *testing.T) { + t.Parallel() + + mockLogger := &mockStyledLogger{} + + trans := &mockTranslator{ + name: "anthropic", + implementsTokenCounter: true, + implementsTokenSerializer: true, + countTokensFunc: func(_ context.Context, _ *http.Request) (*translator.TokenCountResponse, error) { + return &translator.TokenCountResponse{InputTokens: 17}, nil + }, + serialiseCountTokensFunc: func(resp *translator.TokenCountResponse) ([]byte, error) { + return json.Marshal(map[string]int{"input_tokens": resp.InputTokens}) + }, + } + + app := &Application{logger: mockLogger} + handler := app.tokenCountHandler(trans) + + reqBody, _ := json.Marshal(map[string]interface{}{ + "model": "test-model", + "messages": []interface{}{map[string]interface{}{"role": "user", "content": "hi"}}, + }) + req := httptest.NewRequest(http.MethodPost, "/olla/anthropic/v1/messages/count_tokens", bytes.NewReader(reqBody)) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + var body map[string]int + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body)) + require.Equal(t, 17, body["input_tokens"]) +} + +// TestCommittedResponseWriter_WriteAndWriteHeaderSetCommitted confirms the existing +// committed-flag behaviour for Write and WriteHeader is unaffected by the new methods. +func TestCommittedResponseWriter_WriteAndWriteHeaderSetCommitted(t *testing.T) { + t.Parallel() + + t.Run("Write", func(t *testing.T) { + t.Parallel() + rec := httptest.NewRecorder() + cw := newCommittedResponseWriter(rec) + _, err := cw.Write([]byte("hello")) + require.NoError(t, err) + if !cw.committed.Load() { + t.Fatal("committed must be true after Write()") + } + }) + + t.Run("WriteHeader", func(t *testing.T) { + t.Parallel() + rec := httptest.NewRecorder() + cw := newCommittedResponseWriter(rec) + cw.WriteHeader(http.StatusOK) + if !cw.committed.Load() { + t.Fatal("committed must be true after WriteHeader()") + } + }) +} diff --git a/internal/app/handlers/handler_translator_models.go b/internal/app/handlers/handler_translator_models.go index bc3f18b9..66a15dbd 100644 --- a/internal/app/handlers/handler_translator_models.go +++ b/internal/app/handlers/handler_translator_models.go @@ -4,6 +4,9 @@ import ( "context" "encoding/json" "net/http" + "sort" + "strings" + "sync" "time" "github.com/thushan/olla/internal/adapter/translator" @@ -11,8 +14,76 @@ import ( "github.com/thushan/olla/internal/core/domain" ) +// processStartTime is captured once at startup so that created_at in the +// Anthropic models response is stable across calls. There is no genuine +// creation timestamp on UnifiedModel (LastSeen is a discovery heartbeat that +// shifts on every health poll), so process start is the best stable proxy. +var processStartTime = time.Now().UTC().Format(time.RFC3339) + +// anthropicModelsCache is a one-entry cache for the Anthropic GET /v1/models +// response. The underlying model set changes only on the 30-second discovery +// cycle, so rebuilding it on every HTTP request is pure waste under load. +// +// Invalidation key: the sorted, comma-joined set of healthy model IDs. A change +// in any model ID (add, remove, rename) produces a different fingerprint and +// triggers a rebuild. Endpoint health changes are captured because we filter +// by healthy endpoints before building the fingerprint. +type anthropicModelsCache struct { + fingerprint string + body []byte // pre-encoded JSON ready to write + mu sync.RWMutex +} + +// get returns the cached body if the fingerprint matches, otherwise nil. +func (c *anthropicModelsCache) get(fingerprint string) []byte { + c.mu.RLock() + defer c.mu.RUnlock() + if c.fingerprint == fingerprint { + return c.body + } + return nil +} + +// set stores a new fingerprint and pre-encoded body. +func (c *anthropicModelsCache) set(fingerprint string, body []byte) { + c.mu.Lock() + c.fingerprint = fingerprint + c.body = body + c.mu.Unlock() +} + +// modelSetFingerprint builds a stable, cheap change-detection key from the +// set of models as they will be emitted to clients. +// +// The response uses the first alias as the emitted id, not UnifiedModel.ID, so +// the fingerprint must include the emitted value. Without this, an alias change +// (e.g. a version suffix update) while the underlying IDs are stable would leave +// stale alias strings in the cache until the next model-set change. Each entry is +// Each entry is "id\x00emittedID" (null-byte separator) so that changes to +// either field invalidate the cache. A colon would collide with IDs that contain +// colons (e.g. "a:b" + alias "c" vs "a" + alias "b:c" both produce "a:b:c"). +// Sorting ensures the key is order-independent. +func modelSetFingerprint(models []*domain.UnifiedModel) string { + if len(models) == 0 { + return "" + } + ids := make([]string, 0, len(models)) + for _, m := range models { + emitted := m.ID + if len(m.Aliases) > 0 { + emitted = m.Aliases[0].Name + } + ids = append(ids, m.ID+"\x00"+emitted) + } + sort.Strings(ids) + return strings.Join(ids, ",") +} + // list models in translator format (eg /olla/anthropic/v1/models) func (a *Application) translatorModelsHandler(trans translator.RequestTranslator) http.HandlerFunc { + // One cache per translator registration; lives for the lifetime of the handler. + cache := &anthropicModelsCache{} + return func(w http.ResponseWriter, r *http.Request) { ctx := r.Context() @@ -28,68 +99,97 @@ func (a *Application) translatorModelsHandler(trans translator.RequestTranslator if getter, ok := a.modelRegistry.(unifiedModelsGetter); ok { unifiedModels, err = getter.GetUnifiedModels(ctx) if err != nil { - a.writeTranslatorModelsError(w, trans, "failed to get unified models", http.StatusInternalServerError) + a.writeTranslatorModelsError(w, trans, "failed to get unified models") return } } else { - a.writeTranslatorModelsError(w, trans, "unified models not supported", http.StatusInternalServerError) + a.writeTranslatorModelsError(w, trans, "unified models not supported") return } // skip unhealthy models for endpoint response healthyModels, err := a.filterModelsByHealth(ctx, unifiedModels) if err != nil { - a.writeTranslatorModelsError(w, trans, "failed to filter models by health", http.StatusInternalServerError) + a.writeTranslatorModelsError(w, trans, "failed to filter models by health") + return + } + + // Cache hit: the model set has not changed since the last build. + // GetUnifiedModels + filterModelsByHealth above are still called on every + // request so we always have fresh health data for the fingerprint, but the + // expensive JSON encoding and map allocation are skipped on a hit. + fp := modelSetFingerprint(healthyModels) + if body := cache.get(fp); body != nil { + w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(body) return } - // Anthropic format matches the Python reference: {data: [{id, name, created, description, type}]} + // Cache miss: build the Anthropic-format response and store it. response := a.convertModelsToAnthropicFormat(healthyModels) + body, err := json.Marshal(response) + if err != nil { + a.writeTranslatorModelsError(w, trans, "failed to encode models response") + return + } + cache.set(fp, body) w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(response) + _, _ = w.Write(body) } } -// convert to anthropic format (matches python reference) +// convertModelsToAnthropicFormat emits the Anthropic GET /v1/models wire format. +// The SDK uses "type" as a deserialise discriminator so it must be "model", not "chat". +// "display_name" and "created_at" (RFC3339) match the published Anthropic spec. func (a *Application) convertModelsToAnthropicFormat(models []*domain.UnifiedModel) map[string]interface{} { data := make([]map[string]interface{}, 0, len(models)) for _, model := range models { - // use first alias or fallback to default id + // prefer the first alias as the canonical id clients send back to us modelID := model.ID if len(model.Aliases) > 0 { modelID = model.Aliases[0].Name } - // map model fields to human-readable fields entry := map[string]interface{}{ - "id": modelID, - "name": modelID, - "created": time.Now().Unix(), - "description": "Chat model via Olla proxy", - "type": "chat", + "type": "model", + "id": modelID, + "display_name": modelID, + "created_at": processStartTime, } data = append(data, entry) } + // first_id / last_id are null when the list is empty + var firstID, lastID interface{} + if len(data) > 0 { + firstID = data[0]["id"] + lastID = data[len(data)-1]["id"] + } + return map[string]interface{}{ - "data": data, + "data": data, + "has_more": false, + "first_id": firstID, + "last_id": lastID, } } -// send error response for models endpoint -func (a *Application) writeTranslatorModelsError(w http.ResponseWriter, trans translator.RequestTranslator, message string, statusCode int) { +// send error response for models endpoint. All failure paths here are 500s — +// either the registry is unavailable or JSON encoding failed, both of which are +// internal failures rather than client errors. +func (a *Application) writeTranslatorModelsError(w http.ResponseWriter, trans translator.RequestTranslator, message string) { a.logger.Error("Translator models request failed", "translator", trans.Name(), - "error", message, - "status", statusCode) + "error", message) // use custom error format or fallback to generic if errorWriter, ok := trans.(translator.ErrorWriter); ok { - errorWriter.WriteError(w, err{message: message}, statusCode) + errorWriter.WriteError(w, err{message: message}, http.StatusInternalServerError) return } @@ -101,7 +201,7 @@ func (a *Application) writeTranslatorModelsError(w http.ResponseWriter, trans tr } w.Header().Set(constants.HeaderContentType, constants.ContentTypeJSON) - w.WriteHeader(statusCode) + w.WriteHeader(http.StatusInternalServerError) json.NewEncoder(w).Encode(errorResp) } diff --git a/internal/app/handlers/handler_translator_models_cache_test.go b/internal/app/handlers/handler_translator_models_cache_test.go new file mode 100644 index 00000000..8593fa77 --- /dev/null +++ b/internal/app/handlers/handler_translator_models_cache_test.go @@ -0,0 +1,338 @@ +package handlers + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/thushan/olla/internal/adapter/translator" + "github.com/thushan/olla/internal/adapter/translator/anthropic" + "github.com/thushan/olla/internal/config" + "github.com/thushan/olla/internal/core/domain" +) + +// TestAnthropicModelsCache_HitReturnsSameContent verifies that a cache hit +// returns byte-identical content to the first (miss) response. +func TestAnthropicModelsCache_HitReturnsSameContent(t *testing.T) { + t.Parallel() + + models := []*domain.UnifiedModel{ + { + ID: "claude/3:opus", + Aliases: []domain.AliasEntry{{Name: "claude-3-opus", Source: "test"}}, + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }, + { + ID: "claude/3:sonnet", + Aliases: []domain.AliasEntry{{Name: "claude-3-sonnet", Source: "test"}}, + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + modelRegistry: &mockTranslatorRegistry{models: models}, + discoveryService: &mockDiscoveryService{}, + repository: &mockEndpointRepository{}, + } + + trans := mustNewAnthropicTranslator(t) + handler := app.translatorModelsHandler(trans) + + body1 := doGetModels(t, handler) + body2 := doGetModels(t, handler) + + if string(body1) != string(body2) { + t.Errorf("cache hit returned different content:\n miss: %s\n hit: %s", body1, body2) + } + + // Both responses must be valid JSON with the expected envelope. + validateModelsEnvelope(t, body1) +} + +// TestAnthropicModelsCache_InvalidatesOnModelSetChange verifies that when the +// set of healthy models changes, the cache is invalidated and the response is +// rebuilt with the updated data. +func TestAnthropicModelsCache_InvalidatesOnModelSetChange(t *testing.T) { + t.Parallel() + + initialModels := []*domain.UnifiedModel{ + { + ID: "model/a", + Aliases: []domain.AliasEntry{{Name: "model-a", Source: "test"}}, + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }, + } + + reg := &mutableTranslatorRegistry{models: initialModels} + app := &Application{ + logger: &mockStyledLogger{}, + modelRegistry: reg, + discoveryService: &mockDiscoveryService{}, + repository: &mockEndpointRepository{}, + } + + trans := mustNewAnthropicTranslator(t) + handler := app.translatorModelsHandler(trans) + + // First call: miss, builds cache. + body1 := doGetModels(t, handler) + resp1 := decodeModelsData(t, body1) + if len(resp1) != 1 { + t.Fatalf("expected 1 model, got %d", len(resp1)) + } + + // Simulate a discovery cycle adding a new model. + reg.mu.Lock() + reg.models = append(reg.models, &domain.UnifiedModel{ + ID: "model/b", + Aliases: []domain.AliasEntry{{Name: "model-b", Source: "test"}}, + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }) + reg.mu.Unlock() + + // Second call: fingerprint changed, must rebuild. + body2 := doGetModels(t, handler) + resp2 := decodeModelsData(t, body2) + if len(resp2) != 2 { + t.Fatalf("expected 2 models after registry update, got %d", len(resp2)) + } +} + +// TestAnthropicModelsCache_ConcurrentSafety exercises concurrent reads and a +// simultaneous registry update. The -race flag enforces there are no data races. +func TestAnthropicModelsCache_ConcurrentSafety(t *testing.T) { + if testing.Short() { + t.Skip("skipping concurrency test in short mode") + } + t.Parallel() + + initialModels := []*domain.UnifiedModel{ + { + ID: "concurrent/model", + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }, + } + + reg := &mutableTranslatorRegistry{models: initialModels} + app := &Application{ + logger: &mockStyledLogger{}, + modelRegistry: reg, + discoveryService: &mockDiscoveryService{}, + repository: &mockEndpointRepository{}, + } + + trans := mustNewAnthropicTranslator(t) + handler := app.translatorModelsHandler(trans) + + const readers = 20 + var wg sync.WaitGroup + + for i := range readers { + wg.Add(1) + go func(n int) { + defer wg.Done() + for range 50 { + doGetModels(t, handler) + } + _ = n + }(i) + } + + // Simulate discovery updates while readers are running. + wg.Add(1) + go func() { + defer wg.Done() + for i := range 5 { + reg.mu.Lock() + reg.models = append(reg.models, &domain.UnifiedModel{ + ID: "extra/model-" + string(rune('a'+i)), + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }) + reg.mu.Unlock() + } + }() + + wg.Wait() +} + +// TestModelSetFingerprint verifies the fingerprint is order-independent and +// changes when models are added or removed. +func TestModelSetFingerprint(t *testing.T) { + t.Parallel() + + a := []*domain.UnifiedModel{{ID: "a"}, {ID: "b"}, {ID: "c"}} + b := []*domain.UnifiedModel{{ID: "c"}, {ID: "a"}, {ID: "b"}} // different order + c := []*domain.UnifiedModel{{ID: "a"}, {ID: "b"}} // one fewer + + if modelSetFingerprint(a) != modelSetFingerprint(b) { + t.Error("fingerprint must be order-independent") + } + if modelSetFingerprint(a) == modelSetFingerprint(c) { + t.Error("fingerprint must differ when model set changes") + } + if modelSetFingerprint(nil) != "" { + t.Error("empty model set must produce empty fingerprint") + } +} + +// TestModelSetFingerprint_AliasChange verifies that the fingerprint changes when an +// alias changes even if the underlying model IDs are stable. Without this, the cache +// would serve stale emitted IDs (first alias) for the lifetime of the unchanged ID set. +func TestModelSetFingerprint_AliasChange(t *testing.T) { + t.Parallel() + + before := []*domain.UnifiedModel{ + { + ID: "model/internal", + Aliases: []domain.AliasEntry{{Name: "my-model-v1", Source: "test"}}, + }, + } + after := []*domain.UnifiedModel{ + { + ID: "model/internal", // same underlying ID + Aliases: []domain.AliasEntry{{Name: "my-model-v2", Source: "test"}}, // alias changed + }, + } + + if modelSetFingerprint(before) == modelSetFingerprint(after) { + t.Error("fingerprint must differ when the first alias (emitted id) changes") + } +} + +// TestAnthropicModelsCache_InvalidatesOnAliasChange verifies that when only an alias +// changes (underlying ID stable), the cache is invalidated and the new alias is served. +func TestAnthropicModelsCache_InvalidatesOnAliasChange(t *testing.T) { + t.Parallel() + + reg := &mutableTranslatorRegistry{ + models: []*domain.UnifiedModel{ + { + ID: "internal/model", + Aliases: []domain.AliasEntry{{Name: "model-v1", Source: "test"}}, + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + }, + }, + } + + app := &Application{ + logger: &mockStyledLogger{}, + modelRegistry: reg, + discoveryService: &mockDiscoveryService{}, + repository: &mockEndpointRepository{}, + } + + trans := mustNewAnthropicTranslator(t) + handler := app.translatorModelsHandler(trans) + + // First call: cache miss, builds with model-v1. + body1 := doGetModels(t, handler) + data1 := decodeModelsData(t, body1) + if len(data1) != 1 { + t.Fatalf("expected 1 model, got %d", len(data1)) + } + firstID1, ok := data1[0].(map[string]interface{})["id"].(string) + if !ok { + t.Fatal("expected id to be a string") + } + if firstID1 != "model-v1" { + t.Errorf("expected emitted id %q, got %q", "model-v1", firstID1) + } + + // Simulate a discovery cycle that updates only the alias (same underlying ID). + reg.mu.Lock() + reg.models[0] = &domain.UnifiedModel{ + ID: "internal/model", // unchanged + Aliases: []domain.AliasEntry{{Name: "model-v2", Source: "test"}}, + SourceEndpoints: []domain.SourceEndpoint{{EndpointURL: "http://localhost:8080"}}, + } + reg.mu.Unlock() + + // Second call: fingerprint changed due to alias change, must rebuild. + body2 := doGetModels(t, handler) + data2 := decodeModelsData(t, body2) + if len(data2) != 1 { + t.Fatalf("expected 1 model after alias update, got %d", len(data2)) + } + firstID2, ok := data2[0].(map[string]interface{})["id"].(string) + if !ok { + t.Fatal("expected id to be a string") + } + if firstID2 != "model-v2" { + t.Errorf("cache not invalidated on alias change: expected %q, got %q", "model-v2", firstID2) + } +} + +// mustNewAnthropicTranslator creates an Anthropic translator for test use. +func mustNewAnthropicTranslator(t *testing.T) translator.RequestTranslator { + t.Helper() + cfg := config.AnthropicTranslatorConfig{Enabled: true, MaxMessageSize: 10 << 20} + trans, err := anthropic.NewTranslator(&mockStyledLogger{}, cfg) + if err != nil { + t.Fatalf("failed to create anthropic translator: %v", err) + } + return trans +} + +func doGetModels(t *testing.T, handler http.HandlerFunc) []byte { + t.Helper() + req := httptest.NewRequest(http.MethodGet, "/olla/anthropic/v1/models", nil) + rr := httptest.NewRecorder() + handler(rr, req) + if rr.Code != http.StatusOK { + // t.Errorf rather than t.Fatalf: this helper is called from goroutines + // in the concurrency test, and t.Fatalf calls runtime.Goexit which only + // exits the calling goroutine, silently leaving the parent test passing. + t.Errorf("expected 200, got %d: %s", rr.Code, rr.Body.String()) + return nil + } + // Return a copy so the recorder's buffer can be reused. + out := make([]byte, rr.Body.Len()) + copy(out, rr.Body.Bytes()) + return out +} + +func decodeModelsData(t *testing.T, body []byte) []interface{} { + t.Helper() + var resp map[string]interface{} + if err := json.Unmarshal(body, &resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + data, ok := resp["data"].([]interface{}) + if !ok { + t.Fatalf("expected data to be an array, got %T", resp["data"]) + } + return data +} + +func validateModelsEnvelope(t *testing.T, body []byte) { + t.Helper() + var resp map[string]interface{} + if err := json.Unmarshal(body, &resp); err != nil { + t.Fatalf("invalid JSON: %v", err) + } + for _, key := range []string{"data", "has_more", "first_id", "last_id"} { + if _, ok := resp[key]; !ok { + t.Errorf("missing envelope field %q", key) + } + } +} + +// mutableTranslatorRegistry allows tests to swap the model list mid-test to +// simulate a discovery cycle update. +type mutableTranslatorRegistry struct { + baseMockRegistry + mu sync.RWMutex + models []*domain.UnifiedModel +} + +func (m *mutableTranslatorRegistry) GetUnifiedModels(_ context.Context) ([]*domain.UnifiedModel, error) { + m.mu.RLock() + defer m.mu.RUnlock() + out := make([]*domain.UnifiedModel, len(m.models)) + copy(out, m.models) + return out, nil +} diff --git a/internal/app/handlers/handler_translator_models_test.go b/internal/app/handlers/handler_translator_models_test.go index 97b05cb8..7f80797c 100644 --- a/internal/app/handlers/handler_translator_models_test.go +++ b/internal/app/handlers/handler_translator_models_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/thushan/olla/internal/adapter/translator/anthropic" "github.com/thushan/olla/internal/config" @@ -49,7 +50,10 @@ func TestTranslatorModelsHandler_Success(t *testing.T) { Enabled: true, MaxMessageSize: 10 << 20, // 10MB } - trans := anthropic.NewTranslator(app.logger, testConfig) + trans, err := anthropic.NewTranslator(app.logger, testConfig) + if err != nil { + t.Fatalf("failed to create translator: %v", err) + } // Create request req := httptest.NewRequest(http.MethodGet, "/olla/anthropic/v1/models", nil) @@ -81,22 +85,69 @@ func TestTranslatorModelsHandler_Success(t *testing.T) { t.Errorf("expected 2 models, got %d", len(data)) } - // Verify model format matches Anthropic spec + // Verify the response carries Anthropic-compliant envelope fields. + if _, ok := response["has_more"]; !ok { + t.Error("expected has_more field in response envelope") + } + if _, ok := response["first_id"]; !ok { + t.Error("expected first_id field in response envelope") + } + if _, ok := response["last_id"]; !ok { + t.Error("expected last_id field in response envelope") + } + + // Verify per-entry shape matches the Anthropic spec wire format. if len(data) > 0 { model := data[0].(map[string]interface{}) - requiredFields := []string{"id", "name", "created", "description", "type"} - for _, field := range requiredFields { - if _, ok := model[field]; !ok { - t.Errorf("expected field %s in model response", field) - } + + // "type" must be "model" — the Anthropic SDK uses it as a deserialise discriminator. + if modelType, ok := model["type"].(string); !ok || modelType != "model" { + t.Errorf("expected type \"model\", got %q", model["type"]) } - // Verify type is "chat" - if modelType, ok := model["type"].(string); ok { - if modelType != "chat" { - t.Errorf("expected type 'chat', got '%s'", modelType) + // "display_name" replaces "name". + if _, ok := model["display_name"]; !ok { + t.Error("expected display_name field") + } + if _, present := model["name"]; present { + t.Error("unexpected legacy name field; Anthropic spec uses display_name") + } + + // "created_at" must be an RFC3339 string, not a unix int. + createdAt, ok := model["created_at"].(string) + if !ok { + t.Errorf("expected created_at to be a string, got %T", model["created_at"]) + } else if _, parseErr := time.Parse(time.RFC3339, createdAt); parseErr != nil { + t.Errorf("created_at %q is not valid RFC3339: %v", createdAt, parseErr) + } + if _, present := model["created"]; present { + t.Error("unexpected legacy created field; Anthropic spec uses created_at (RFC3339 string)") + } + + // created_at must be stable across repeated calls — it must not re-evaluate + // time.Now() on every request, which would make caching and idempotency impossible. + req2 := httptest.NewRequest(http.MethodGet, "/olla/anthropic/v1/models", nil) + w2 := httptest.NewRecorder() + handler(w2, req2) + + var response2 map[string]interface{} + if err2 := json.NewDecoder(w2.Body).Decode(&response2); err2 != nil { + t.Fatalf("failed to decode second response: %v", err2) + } + if data2, ok2 := response2["data"].([]interface{}); ok2 && len(data2) > 0 { + model2 := data2[0].(map[string]interface{}) + createdAt2, ok2 := model2["created_at"].(string) + if !ok2 { + t.Errorf("expected created_at to be a string on second call, got %T", model2["created_at"]) + } else if createdAt2 != createdAt { + t.Errorf("created_at changed between calls: %q → %q; must be stable", createdAt, createdAt2) } } + + // "description" is not in the Anthropic spec. + if _, present := model["description"]; present { + t.Error("unexpected description field; not in Anthropic spec") + } } } @@ -109,7 +160,10 @@ func TestTranslatorModelsHandler_EmptyRegistry(t *testing.T) { Enabled: true, MaxMessageSize: 10 << 20, // 10MB } - trans := anthropic.NewTranslator(app.logger, testConfig) + trans, err := anthropic.NewTranslator(app.logger, testConfig) + if err != nil { + t.Fatalf("failed to create translator: %v", err) + } // Create request req := httptest.NewRequest(http.MethodGet, "/olla/anthropic/v1/models", nil) @@ -169,7 +223,10 @@ func TestTranslatorModelsHandler_ResponseFormat(t *testing.T) { Enabled: true, MaxMessageSize: 10 << 20, // 10MB } - trans := anthropic.NewTranslator(app.logger, testConfig) + trans, err := anthropic.NewTranslator(app.logger, testConfig) + if err != nil { + t.Fatalf("failed to create translator: %v", err) + } // Create request req := httptest.NewRequest(http.MethodGet, "/olla/anthropic/v1/models", nil) @@ -185,8 +242,7 @@ func TestTranslatorModelsHandler_ResponseFormat(t *testing.T) { t.Fatalf("failed to decode response: %v", err) } - // Verify structure matches Python reference - // {data: [{id, name, created, description, type}]} + // Verify the Anthropic GET /v1/models wire format. data, ok := response["data"].([]interface{}) if !ok { t.Fatal("expected data field") @@ -196,20 +252,39 @@ func TestTranslatorModelsHandler_ResponseFormat(t *testing.T) { t.Fatal("expected at least one model") } + // Envelope fields. + if hasMo, ok := response["has_more"].(bool); !ok || hasMo { + t.Errorf("expected has_more to be false, got %v", response["has_more"]) + } + if response["first_id"] == nil { + t.Error("expected first_id to be non-null when data is non-empty") + } + if response["last_id"] == nil { + t.Error("expected last_id to be non-null when data is non-empty") + } + model := data[0].(map[string]interface{}) - // Verify model has correct fields if _, ok := model["id"]; !ok { t.Error("expected id field") } - // Verify created is unix timestamp - if created, ok := model["created"].(float64); ok { - if created <= 0 { - t.Error("expected created to be positive unix timestamp") - } - } else { - t.Error("expected created to be number") + // "type" must be "model" per spec. + if modelType, ok := model["type"].(string); !ok || modelType != "model" { + t.Errorf("expected type \"model\", got %q", model["type"]) + } + + // "display_name" replaces legacy "name". + if _, ok := model["display_name"]; !ok { + t.Error("expected display_name field") + } + + // "created_at" must be a valid RFC3339 string. + createdAt, ok := model["created_at"].(string) + if !ok { + t.Errorf("expected created_at to be a string, got %T", model["created_at"]) + } else if _, parseErr := time.Parse(time.RFC3339, createdAt); parseErr != nil { + t.Errorf("created_at %q is not valid RFC3339: %v", createdAt, parseErr) } } diff --git a/internal/app/middleware/logging.go b/internal/app/middleware/logging.go index ceb50be2..6fb2efe9 100644 --- a/internal/app/middleware/logging.go +++ b/internal/app/middleware/logging.go @@ -127,7 +127,9 @@ func EnhancedLoggingMiddleware(styledLogger logger.StyledLogger) func(http.Handl // Add to context for propagation ctx := context.WithValue(r.Context(), RequestIDKey, requestID) - // Create a base logger with request ID + // Create a base logger with request ID. slog.With allocates; it runs + // regardless of level because the logger is stored in context for handlers + // that may log at any level. This is unavoidable and is not gated. baseLogger := slog.Default().With(constants.ContextRequestIdKey, requestID) ctx = context.WithValue(ctx, LoggerKey, baseLogger) @@ -137,47 +139,66 @@ func EnhancedLoggingMiddleware(styledLogger logger.StyledLogger) func(http.Handl // Wrap response writer to capture metrics wrapped := &responseWriter{ResponseWriter: w, status: 200} - // Log request start - use Debug for proxy requests to reduce noise - // Proxy requests will log their own "Request received" at INFO level - logFields := []any{ - "method", r.Method, - "path", r.URL.Path, - "remote_addr", r.RemoteAddr, - "user_agent", r.UserAgent(), - "request_bytes", requestSize, - "request_size_formatted", formatBytes(requestSize), - } - - if IsProxyRequest(r.URL.Path) { - // For proxy requests, just log at debug level since handler will log INFO - baseLogger.Debug("HTTP request started", logFields...) + // Gate field construction on whether the record will actually be emitted. + // On the proxy hot path at the default info level, Debug records are discarded + // by the handler — building the []any slice and calling formatBytes 2x per + // request is pure waste. Non-proxy requests log at Info, so they only pay the + // cost when info-level logging is actually enabled. + isProxy := IsProxyRequest(r.URL.Path) + if isProxy { + if baseLogger.Enabled(ctx, slog.LevelDebug) { + baseLogger.Debug("HTTP request started", + "method", r.Method, + "path", r.URL.Path, + "remote_addr", r.RemoteAddr, + "user_agent", r.UserAgent(), + "request_bytes", requestSize, + "request_size_formatted", formatBytes(requestSize), + ) + } } else { - // For non-proxy requests (health, status, etc), log at INFO - baseLogger.Info("Request started", logFields...) + if baseLogger.Enabled(ctx, slog.LevelInfo) { + baseLogger.Info("Request started", + "method", r.Method, + "path", r.URL.Path, + "remote_addr", r.RemoteAddr, + "user_agent", r.UserAgent(), + "request_bytes", requestSize, + "request_size_formatted", formatBytes(requestSize), + ) + } } next.ServeHTTP(wrapped, r.WithContext(ctx)) duration := time.Since(start) - // Log request completion - use Debug for proxy requests to reduce noise - completionFields := []any{ - "method", r.Method, - "path", r.URL.Path, - "status", wrapped.status, - "duration_ms", duration.Milliseconds(), - "duration_formatted", duration.String(), - "request_bytes", requestSize, - "response_bytes", wrapped.size, - "size_flow", fmt.Sprintf("%s -> %s", formatBytes(requestSize), formatBytes(wrapped.size)), - } - - if IsProxyRequest(r.URL.Path) { - // For proxy requests, just log at debug level since handler will log INFO - baseLogger.Debug("HTTP request completed", completionFields...) + if isProxy { + if baseLogger.Enabled(ctx, slog.LevelDebug) { + baseLogger.Debug("HTTP request completed", + "method", r.Method, + "path", r.URL.Path, + "status", wrapped.status, + "duration_ms", duration.Milliseconds(), + "duration_formatted", duration.String(), + "request_bytes", requestSize, + "response_bytes", wrapped.size, + "size_flow", fmt.Sprintf("%s -> %s", formatBytes(requestSize), formatBytes(wrapped.size)), + ) + } } else { - // For non-proxy requests, log at INFO - baseLogger.Info("Request completed", completionFields...) + if baseLogger.Enabled(ctx, slog.LevelInfo) { + baseLogger.Info("Request completed", + "method", r.Method, + "path", r.URL.Path, + "status", wrapped.status, + "duration_ms", duration.Milliseconds(), + "duration_formatted", duration.String(), + "request_bytes", requestSize, + "response_bytes", wrapped.size, + "size_flow", fmt.Sprintf("%s -> %s", formatBytes(requestSize), formatBytes(wrapped.size)), + ) + } } }) } @@ -209,26 +230,29 @@ func AccessLoggingMiddleware(styledLogger logger.StyledLogger) func(http.Handler duration := time.Since(start) - // Create detailed context for file logging only - detailedCtx := context.WithValue(r.Context(), logger.DefaultDetailedCookie, true) - - // Log detailed access information (file only) + // Access logs route to a dedicated file handler (keyed by DefaultDetailedCookie). + // Build fields only when the handler is enabled to avoid allocating the + // format.RFC3339 string, redactQuery output, and the variadic slice on every + // request when the file sink is not configured. baseLogger := slog.Default() - baseLogger.InfoContext(detailedCtx, "Access log", - "timestamp", start.Format(time.RFC3339), - "request_id", requestID, - "remote_addr", r.RemoteAddr, - "method", r.Method, - "path", r.URL.Path, - "query", redactQuery(r.URL.RawQuery), - "status", wrapped.status, - "request_bytes", requestSize, - "response_bytes", wrapped.size, - "duration_ms", duration.Milliseconds(), - "user_agent", r.UserAgent(), - "referer", r.Referer(), - "content_type", r.Header.Get(constants.HeaderContentType), - "accept", r.Header.Get(constants.HeaderAccept)) + detailedCtx := context.WithValue(r.Context(), logger.DefaultDetailedCookie, true) + if baseLogger.Enabled(detailedCtx, slog.LevelInfo) { + baseLogger.InfoContext(detailedCtx, "Access log", + "timestamp", start.Format(time.RFC3339), + "request_id", requestID, + "remote_addr", r.RemoteAddr, + "method", r.Method, + "path", r.URL.Path, + "query", redactQuery(r.URL.RawQuery), + "status", wrapped.status, + "request_bytes", requestSize, + "response_bytes", wrapped.size, + "duration_ms", duration.Milliseconds(), + "user_agent", r.UserAgent(), + "referer", r.Referer(), + "content_type", r.Header.Get(constants.HeaderContentType), + "accept", r.Header.Get(constants.HeaderAccept)) + } }) } } diff --git a/internal/app/middleware/logging_gate_test.go b/internal/app/middleware/logging_gate_test.go new file mode 100644 index 00000000..43f6240f --- /dev/null +++ b/internal/app/middleware/logging_gate_test.go @@ -0,0 +1,145 @@ +package middleware + +import ( + "context" + "log/slog" + "net/http" + "net/http/httptest" + "testing" +) + +// countingHandler counts Handle calls, discards output, and respects a fixed min level. +// Used to assert whether log records are emitted without depending on stdout output. +type countingHandler struct { + minLevel slog.Level + count int +} + +func (h *countingHandler) Enabled(_ context.Context, level slog.Level) bool { + return level >= h.minLevel +} +func (h *countingHandler) Handle(_ context.Context, _ slog.Record) error { + h.count++ + return nil +} +func (h *countingHandler) WithAttrs(_ []slog.Attr) slog.Handler { return h } +func (h *countingHandler) WithGroup(_ string) slog.Handler { return h } + +// setDefaultLogger replaces slog.Default for the duration of the test. +// NOT safe to call from t.Parallel() tests as slog.Default is process-global; +// these tests must run serially via subtests under a single parent. +func setDefaultLogger(t *testing.T, h slog.Handler) { + t.Helper() + prev := slog.Default() + slog.SetDefault(slog.New(h)) + t.Cleanup(func() { slog.SetDefault(prev) }) +} + +// TestEnhancedLogging_Gate exercises the level-gate logic in EnhancedLoggingMiddleware. +// All subtests share the same slog.Default mutation so they must run serially. +func TestEnhancedLogging_Gate(t *testing.T) { + mockLogger := &mockStyledLogger{} + noop := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + // ProxyPath_InfoLevel: at info level, Debug records must not be emitted. + // The responseWriter wrap and request-ID propagation must still occur. + t.Run("ProxyPath_InfoLevel", func(t *testing.T) { + ch := &countingHandler{minLevel: slog.LevelInfo} + setDefaultLogger(t, ch) + + var gotRequestID string + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotRequestID = GetRequestID(r.Context()) + w.WriteHeader(http.StatusOK) + }) + mw := EnhancedLoggingMiddleware(mockLogger)(inner) + + req := httptest.NewRequest(http.MethodPost, "/olla/ollama/api/chat", nil) + rr := httptest.NewRecorder() + mw.ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rr.Code) + } + // The request-ID header must still be set - it is outside the log gate. + if rr.Header().Get("X-Olla-Request-ID") == "" { + t.Error("X-Olla-Request-ID header must be set even when debug logging is suppressed") + } + // The context request ID must still be propagated for downstream handlers. + if gotRequestID == "" { + t.Error("request ID must be in context even when debug logging is suppressed") + } + // "HTTP request started" and "HTTP request completed" are both Debug. + // With minLevel=Info the gate must suppress both records. + if ch.count != 0 { + t.Errorf("expected 0 log records at info level for proxy path, got %d", ch.count) + } + }) + + // ProxyPath_DebugLevel: at debug level, both records must be emitted. + t.Run("ProxyPath_DebugLevel", func(t *testing.T) { + ch := &countingHandler{minLevel: slog.LevelDebug} + setDefaultLogger(t, ch) + + mw := EnhancedLoggingMiddleware(mockLogger)(noop) + + req := httptest.NewRequest(http.MethodPost, "/olla/ollama/api/chat", nil) + rr := httptest.NewRecorder() + mw.ServeHTTP(rr, req) + + // Lower bound of 2: start + completion. Handler may emit more via slog.With. + if ch.count < 2 { + t.Errorf("expected at least 2 log records at debug level for proxy path, got %d", ch.count) + } + }) + + // NonProxyPath_InfoLevel: non-proxy requests log at Info, which IS enabled at + // the default level - both "Request started" and "Request completed" must appear. + t.Run("NonProxyPath_InfoLevel", func(t *testing.T) { + ch := &countingHandler{minLevel: slog.LevelInfo} + setDefaultLogger(t, ch) + + mw := EnhancedLoggingMiddleware(mockLogger)(noop) + + req := httptest.NewRequest(http.MethodGet, "/internal/health", nil) + rr := httptest.NewRecorder() + mw.ServeHTTP(rr, req) + + if ch.count < 2 { + t.Errorf("expected at least 2 log records at info level for non-proxy path, got %d", ch.count) + } + }) +} + +// BenchmarkEnhancedLogging_ProxyPath_InfoLevel measures the hot-path overhead +// of EnhancedLoggingMiddleware at the default info level for proxy requests. +// At info level all Debug records are suppressed, so formatBytes, the []any field +// slices, and fmt.Sprintf must be skipped entirely. +// +// Run with: +// +// go test -bench=BenchmarkEnhancedLogging_ProxyPath_InfoLevel -benchmem ./internal/app/middleware/ +func BenchmarkEnhancedLogging_ProxyPath_InfoLevel(b *testing.B) { + ch := &countingHandler{minLevel: slog.LevelInfo} + prev := slog.Default() + slog.SetDefault(slog.New(ch)) + b.Cleanup(func() { slog.SetDefault(prev) }) + + mockLogger := &mockStyledLogger{} + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + mw := EnhancedLoggingMiddleware(mockLogger)(inner) + + b.ResetTimer() + b.ReportAllocs() + + for range b.N { + req := httptest.NewRequest(http.MethodPost, "/olla/ollama/api/chat", nil) + rr := httptest.NewRecorder() + mw.ServeHTTP(rr, req) + } +} diff --git a/internal/app/services/http.go b/internal/app/services/http.go index fec08fb6..a823b25e 100644 --- a/internal/app/services/http.go +++ b/internal/app/services/http.go @@ -4,11 +4,10 @@ import ( "context" "errors" "fmt" + "net" "net/http" "time" - "github.com/thushan/olla/internal/util" - "github.com/thushan/olla/internal/app/handlers" "github.com/thushan/olla/internal/app/middleware" "github.com/thushan/olla/internal/config" @@ -67,10 +66,6 @@ func (s *HTTPService) Name() string { func (s *HTTPService) Start(ctx context.Context) error { s.logger.Info("Initialising HTTP service") - if !util.IsPortAvailable(s.config.Host, s.config.Port) { - return fmt.Errorf("port %d is already in use on host %s", s.config.Port, s.config.Host) - } - // Resolve service dependencies now that all services are started if s.statsService != nil { collector, err := s.statsService.GetCollector() @@ -137,7 +132,7 @@ func (s *HTTPService) Start(ctx context.Context) error { } s.application = app - // Wire sticky session stats if enabled — proxySvc holds the wrapper. + // Wire sticky session stats if enabled - proxySvc holds the wrapper. if s.proxySvc != nil { s.application.SetStickyStatsFn(s.proxySvc.StickyStats) } @@ -146,8 +141,8 @@ func (s *HTTPService) Start(ctx context.Context) error { // This is separate from the security chain (which handles proxy routes) so // non-proxy routes are protected without requiring full chain enforcement. if s.securitySvc != nil { - if realAdapters, err := s.securitySvc.GetAdapters(); err == nil && realAdapters != nil { - s.application.SetSecurityAdapters(realAdapters) + if adapters, adapterErr := s.securitySvc.GetAdapters(); adapterErr == nil && adapters != nil { + s.application.SetSecurityAdapters(adapters) } } @@ -168,8 +163,15 @@ func (s *HTTPService) Start(ctx context.Context) error { "allow_credentials", s.fullConfig.Server.Cors.AllowCredentials) } + addr := s.config.GetAddress() + + ln, err := bindListener(ctx, addr) + if err != nil { + return err + } + s.server = &http.Server{ - Addr: s.config.GetAddress(), + Addr: addr, Handler: root, ReadTimeout: readTimeout, ReadHeaderTimeout: readHeaderTimeout, @@ -177,22 +179,20 @@ func (s *HTTPService) Start(ctx context.Context) error { IdleTimeout: idleTimeout, } + s.logger.Info("HTTP server listening", + "address", addr, + "readTimeout", readTimeout, + "readHeaderTimeout", readHeaderTimeout, + "writeTimeout", writeTimeout, + "idleTimeout", idleTimeout) + go func() { - s.logger.Info("HTTP server listening", - "address", s.server.Addr, - "readTimeout", readTimeout, - "readHeaderTimeout", readHeaderTimeout, - "writeTimeout", writeTimeout, - "idleTimeout", idleTimeout) - - if serr := s.server.ListenAndServe(); serr != nil && !errors.Is(serr, http.ErrServerClosed) { + if serr := s.server.Serve(ln); serr != nil && !errors.Is(serr, http.ErrServerClosed) { s.logger.Error("HTTP server error", "error", serr) } }() - time.Sleep(100 * time.Millisecond) - - s.logger.Info("Olla started, waiting for requests...", "bind", s.server.Addr) + s.logger.Info("Olla started, waiting for requests...", "bind", addr) s.printWarnings() return nil @@ -214,6 +214,17 @@ func applyCORS(handler http.Handler, cfg config.CorsConfig) http.Handler { return middleware.NewCORS(cfg).Handler(handler) } +// bindListener binds a TCP listener on addr synchronously so that any +// port-in-use error surfaces immediately to the caller. Using ListenConfig +// means the listener also respects context cancellation during startup. +func bindListener(ctx context.Context, addr string) (net.Listener, error) { + ln, err := (&net.ListenConfig{}).Listen(ctx, "tcp", addr) + if err != nil { + return nil, fmt.Errorf("failed to bind %s: %w", addr, err) + } + return ln, nil +} + func (s *HTTPService) printWarnings() { if s.fullConfig.Translators.Anthropic.Inspector.Enabled { s.logger.Warn("Anthropic Inspector is ENABLED. DO NOT USE IN PROD.", "more_info", "https://thushan.github.io/olla/notes/anthropic-inspector/") diff --git a/internal/app/services/http_bind_test.go b/internal/app/services/http_bind_test.go new file mode 100644 index 00000000..b1386b5e --- /dev/null +++ b/internal/app/services/http_bind_test.go @@ -0,0 +1,41 @@ +package services + +import ( + "context" + "net" + "testing" +) + +// TestBindListener_FreePort verifies that bindListener returns a valid listener +// when no other process holds the port. +func TestBindListener_FreePort(t *testing.T) { + t.Parallel() + + ln, err := bindListener(context.Background(), "127.0.0.1:0") + if err != nil { + t.Fatalf("expected bind to succeed on a free port: %v", err) + } + ln.Close() +} + +// TestBindListener_OccupiedPort verifies that bindListener returns a non-nil +// error when another listener already holds the port. This is the path that +// Start() uses to surface port conflicts immediately rather than swallowing +// them inside a goroutine. +func TestBindListener_OccupiedPort(t *testing.T) { + t.Parallel() + + // Hold a port so the second bind must fail. + anchor, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to acquire anchor listener: %v", err) + } + defer anchor.Close() + + addr := anchor.Addr().String() + + _, bindErr := bindListener(context.Background(), addr) + if bindErr == nil { + t.Fatal("expected bindListener to return an error while anchor holds the port, but it succeeded") + } +} diff --git a/internal/app/services/proxy.go b/internal/app/services/proxy.go index 18bc3be7..74d06de9 100644 --- a/internal/app/services/proxy.go +++ b/internal/app/services/proxy.go @@ -153,6 +153,15 @@ func (s *ProxyServiceWrapper) Start(ctx context.Context) error { return nil } +// cleanupable is an optional interface implemented by proxy engines that own +// background goroutines or connection pools that must be released on shutdown. +// Using a type-assert rather than embedding in ports.ProxyService keeps the +// interface minimal and avoids forcing all current and future implementations +// to add a no-op method. +type cleanupable interface { + Cleanup() +} + // Stop gracefully shuts down the proxy service func (s *ProxyServiceWrapper) Stop(ctx context.Context) error { s.logger.Info(" Stopping proxy service") @@ -161,6 +170,14 @@ func (s *ProxyServiceWrapper) Stop(ctx context.Context) error { s.stickyWrapper.Stop() } + // Call Cleanup on the proxy engine if it supports it. This stops the + // cleanupLoop goroutine and releases per-endpoint connection pools, + // preventing a goroutine/resource leak if the engine is ever rebuilt + // in-process. + if c, ok := s.proxyService.(cleanupable); ok { + c.Cleanup() + } + defer func() { s.logger.ResetLine() s.logger.InfoWithStatus("Stopping proxy service", "OK") diff --git a/internal/app/services/proxy_cleanup_test.go b/internal/app/services/proxy_cleanup_test.go new file mode 100644 index 00000000..2272d89d --- /dev/null +++ b/internal/app/services/proxy_cleanup_test.go @@ -0,0 +1,117 @@ +package services + +import ( + "context" + "net/http" + "sync/atomic" + "testing" + + "github.com/thushan/olla/internal/core/domain" + "github.com/thushan/olla/internal/core/ports" + "github.com/thushan/olla/internal/logger" +) + +func newTestLogger() logger.StyledLogger { + cfg := &logger.Config{Level: "error", Theme: "default"} + log, _, _ := logger.New(cfg) + return logger.NewPlainStyledLogger(log) +} + +// spyProxyService is a minimal ports.ProxyService that also implements the +// cleanupable interface. It counts Cleanup invocations for assertion. +type spyProxyService struct { + cleanupCalls atomic.Int32 +} + +func (s *spyProxyService) ProxyRequest(_ context.Context, _ http.ResponseWriter, _ *http.Request, _ *ports.RequestStats, _ logger.StyledLogger) error { + return nil +} + +func (s *spyProxyService) ProxyRequestToEndpoints(_ context.Context, _ http.ResponseWriter, _ *http.Request, _ []*domain.Endpoint, _ *ports.RequestStats, _ logger.StyledLogger) error { + return nil +} + +func (s *spyProxyService) GetStats(_ context.Context) (ports.ProxyStats, error) { + return ports.ProxyStats{}, nil +} + +func (s *spyProxyService) UpdateConfig(_ ports.ProxyConfiguration) {} + +// Cleanup satisfies the optional cleanupable interface. +func (s *spyProxyService) Cleanup() { + s.cleanupCalls.Add(1) +} + +// TestProxyServiceWrapper_Stop_CallsCleanup asserts that Stop invokes the proxy +// engine's Cleanup when the engine implements cleanupable. +func TestProxyServiceWrapper_Stop_CallsCleanup(t *testing.T) { + t.Parallel() + + spy := &spyProxyService{} + w := &ProxyServiceWrapper{ + logger: newTestLogger(), + proxyService: spy, + } + + if err := w.Stop(context.Background()); err != nil { + t.Fatalf("Stop returned unexpected error: %v", err) + } + if got := spy.cleanupCalls.Load(); got != 1 { + t.Errorf("Cleanup called %d times, want 1", got) + } +} + +// TestProxyServiceWrapper_Stop_DoubleCallDoesNotPanic asserts that calling Stop +// twice does not panic. Engine-level idempotency (sync.Once) is exercised +// separately in the olla service tests. +func TestProxyServiceWrapper_Stop_DoubleCallDoesNotPanic(t *testing.T) { + t.Parallel() + + spy := &spyProxyService{} + w := &ProxyServiceWrapper{ + logger: newTestLogger(), + proxyService: spy, + } + + for range 2 { + if err := w.Stop(context.Background()); err != nil { + t.Fatalf("Stop returned error: %v", err) + } + } + // The wrapper does not deduplicate; each Stop call propagates to the engine. + // The engine's own sync.Once handles its internal idempotency. + if got := spy.cleanupCalls.Load(); got != 2 { + t.Errorf("Cleanup called %d times, want 2", got) + } +} + +// bareProxyService implements ports.ProxyService but NOT cleanupable. +type bareProxyService struct{} + +func (b *bareProxyService) ProxyRequest(_ context.Context, _ http.ResponseWriter, _ *http.Request, _ *ports.RequestStats, _ logger.StyledLogger) error { + return nil +} + +func (b *bareProxyService) ProxyRequestToEndpoints(_ context.Context, _ http.ResponseWriter, _ *http.Request, _ []*domain.Endpoint, _ *ports.RequestStats, _ logger.StyledLogger) error { + return nil +} + +func (b *bareProxyService) GetStats(_ context.Context) (ports.ProxyStats, error) { + return ports.ProxyStats{}, nil +} + +func (b *bareProxyService) UpdateConfig(_ ports.ProxyConfiguration) {} + +// TestProxyServiceWrapper_Stop_NoCleanupable asserts Stop does not panic when the +// proxy engine does not implement cleanupable. +func TestProxyServiceWrapper_Stop_NoCleanupable(t *testing.T) { + t.Parallel() + + w := &ProxyServiceWrapper{ + logger: newTestLogger(), + proxyService: &bareProxyService{}, + } + if err := w.Stop(context.Background()); err != nil { + t.Fatalf("Stop returned unexpected error: %v", err) + } +} diff --git a/internal/core/constants/context.go b/internal/core/constants/context.go index c160b908..f03472cf 100644 --- a/internal/core/constants/context.go +++ b/internal/core/constants/context.go @@ -26,4 +26,9 @@ const ( // when a model alias is resolved, allowing the proxy to rewrite the model name in the // request body to match what the selected backend expects ContextModelAliasMapKey = "model_alias_map" + + // ContextInputTokensKey carries a prompt token estimate through the translation + // pipeline so the streaming translator can populate input_tokens in message_start + // before the upstream usage chunk arrives (which only comes at the end of the stream). + ContextInputTokensKey = contextKey("input-tokens") ) diff --git a/internal/core/domain/profile_config.go b/internal/core/domain/profile_config.go index 46b6c28e..5963379e 100644 --- a/internal/core/domain/profile_config.go +++ b/internal/core/domain/profile_config.go @@ -133,6 +133,15 @@ type ContextPattern struct { Context int64 `yaml:"context"` } +// Limitation name constants for features that may be absent on a given backend. +// The passthrough endpoint filter uses exact-name matching against these values, +// so token-counting limitations (which use different names) are never affected. +const ( + AnthropicLimitationNoToolUse = "no_tool_use" + AnthropicLimitationNoExtendedThinking = "no_extended_thinking" + AnthropicLimitationNoVision = "no_vision" +) + // AnthropicSupportConfig declares native Anthropic Messages API support for a // backend platform. This enables the passthrough optimisation: when a backend // natively understands the Anthropic wire format, requests can be forwarded @@ -145,7 +154,7 @@ type ContextPattern struct { // enabled: true // messages_path: "/v1/messages" // token_count: true -// min_version: "2023-06-01" +// min_version: "b4847" // limitations: // - "no_extended_thinking" // - "max_tokens_4096" @@ -154,15 +163,16 @@ type AnthropicSupportConfig struct { // requests (e.g. "/v1/messages"). Required when Enabled is true. MessagesPath string `yaml:"messages_path"` - // MinVersion is the minimum anthropic-version header value the backend - // requires. If the incoming request specifies an older version, the - // translator falls back to the translation path. Use the standard - // Anthropic version date format (e.g. "2023-06-01"). + // MinVersion is the backend build version that introduced native Anthropic + // support (e.g. "b4847", "0.14.0", "0.4.1"). Informational only; not + // currently enforced at runtime. MinVersion string `yaml:"min_version,omitempty"` // Limitations lists Anthropic features this backend does NOT support. - // Used by CanPassthrough to decide whether a particular request can be - // sent directly or must go through translation instead. + // Enforced by the passthrough endpoint filter for messages-affecting + // limitations (no_tool_use, no_extended_thinking, no_vision). Endpoints + // that declare a limitation matching an active request feature are excluded + // from the passthrough subset and fall back to translation. // Common values: "no_extended_thinking", "no_tool_use", "no_vision", // "max_tokens_4096". Limitations []string `yaml:"limitations,omitempty"` diff --git a/internal/util/request.go b/internal/util/request.go index 5c2e7fb0..10e0db07 100644 --- a/internal/util/request.go +++ b/internal/util/request.go @@ -36,6 +36,16 @@ func GetClientIP(r *http.Request, trustProxyHeaders bool, trustedCIDRs []*net.IP return r.RemoteAddr } + // With no trusted CIDRs configured, every source is implicitly untrusted and + // getSourceIP would allocate a net.IP only to immediately discard it. Skip the + // parse and fall through to RemoteAddr directly. + if len(trustedCIDRs) == 0 { + if ip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { + return ip + } + return r.RemoteAddr + } + sourceIP := getSourceIP(r) if sourceIP == nil || !isIPInTrustedCIDRs(sourceIP, trustedCIDRs) { if ip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil { diff --git a/pkg/eventbus/eventbus_worker.go b/pkg/eventbus/eventbus_worker.go index 3cfb5497..aa9e7817 100644 --- a/pkg/eventbus/eventbus_worker.go +++ b/pkg/eventbus/eventbus_worker.go @@ -70,12 +70,16 @@ func (wp *WorkerPool[T]) worker() { } } -// Shutdown stops all workers +// Shutdown stops all workers. func (wp *WorkerPool[T]) Shutdown() { - // Cancel context to signal workers to stop + // Cancel the context so workers exit their select loops. wp.cancel() - // Wait for all workers to finish processing + // Wait for all workers to drain and exit before returning. wp.wg.Wait() - // Now safe to close the channel after all workers have stopped - close(wp.eventChan) + // Do NOT close eventChan. Workers exit via ctx cancellation, not via a + // range-over-channel, so the close is unnecessary. More importantly, + // PublishAsync checks ctx.Done() and then sends in two separate selects — + // a goroutine that passes the ctx check before Shutdown runs can still + // attempt a send on a closed channel, causing a panic. Leaving the channel + // open is safe: it is GC'd once the WorkerPool itself goes out of scope. } diff --git a/pkg/eventbus/eventbus_worker_test.go b/pkg/eventbus/eventbus_worker_test.go index 43689a8a..ce43db54 100644 --- a/pkg/eventbus/eventbus_worker_test.go +++ b/pkg/eventbus/eventbus_worker_test.go @@ -121,6 +121,41 @@ func TestWorkerPool_HandlesBackpressure(t *testing.T) { } } +// TestWorkerPool_PublishAsyncShutdownRace exercises the TOCTOU window that existed +// between PublishAsync's ctx check and its eventChan send when Shutdown closed the +// channel concurrently. The fix removes close(eventChan) from Shutdown; workers +// exit via ctx cancellation so the close was never necessary. +// Run with -race to verify no data race or send-on-closed-channel panic. +func TestWorkerPool_PublishAsyncShutdownRace(t *testing.T) { + t.Parallel() + + const goroutines = 50 + const iterations = 200 + + for trial := range 5 { + _ = trial + eb := New[int]() + + var wg sync.WaitGroup + + // Hammer PublishAsync from many goroutines. + for g := range goroutines { + wg.Add(1) + go func(id int) { + defer wg.Done() + for i := range iterations { + eb.PublishAsync(id*iterations + i) + } + }(g) + } + + // Shutdown concurrently with the senders. This used to race and could + // panic on close(eventChan) being observed by a concurrent sender. + eb.Shutdown() + wg.Wait() + } +} + // TestWorkerPool_ConcurrentPublishing verifies concurrent publishing works correctly func TestWorkerPool_ConcurrentPublishing(t *testing.T) { eb := New[string]() diff --git a/pkg/format/format.go b/pkg/format/format.go index 7773ad2e..99d476cb 100644 --- a/pkg/format/format.go +++ b/pkg/format/format.go @@ -88,7 +88,7 @@ func Duration2(d time.Duration) string { return fmt.Sprintf("%dh %dm", hours, minutes) } days := hours / 24 - hours %= hours % 24 + hours %= 24 return fmt.Sprintf("%dd %dh", days, hours) } diff --git a/pkg/format/format_test.go b/pkg/format/format_test.go new file mode 100644 index 00000000..73585b3d --- /dev/null +++ b/pkg/format/format_test.go @@ -0,0 +1,42 @@ +package format + +import ( + "testing" + "time" +) + +func TestDuration2(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + d time.Duration + want string + }{ + // Sub-hour values + {"30 seconds", 30 * time.Second, "30s"}, + {"5 minutes", 5 * time.Minute, "5m"}, + {"90 minutes", 90 * time.Minute, "1h 30m"}, + // Sub-day values + {"1 hour", 1 * time.Hour, "1h 0m"}, + {"23 hours", 23 * time.Hour, "23h 0m"}, + // Exact multiples of 24 hours - these triggered the divide-by-zero panic. + {"24 hours", 24 * time.Hour, "1d 0h"}, + {"48 hours", 48 * time.Hour, "2d 0h"}, + {"72 hours", 72 * time.Hour, "3d 0h"}, + // Non-multiples that produced wrong values before the fix. + {"25 hours", 25 * time.Hour, "1d 1h"}, + {"31 hours", 31 * time.Hour, "1d 7h"}, + {"49 hours", 49 * time.Hour, "2d 1h"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := Duration2(tc.d) + if got != tc.want { + t.Errorf("Duration2(%v) = %q, want %q", tc.d, got, tc.want) + } + }) + } +} diff --git a/pkg/pool/lite_pool.go b/pkg/pool/lite_pool.go index b1f3f417..5a48a1a2 100644 --- a/pkg/pool/lite_pool.go +++ b/pkg/pool/lite_pool.go @@ -10,25 +10,30 @@ package pool // type assertion in Get() is safe and explicitly silenced. // // Example: -// type RequestContext struct { ... } -// func (r *RequestContext) Reset() { ... } // -// pool, err := NewLitePool(func() *RequestContext { -// return &RequestContext{} -// }) -// if err != nil { -// // handle error -// } +// type RequestContext struct { ... } +// func (r *RequestContext) Reset() { ... } // -// ctx := pool.Get() -// ... -// pool.Put(ctx) +// pool, err := NewLitePool(func() *RequestContext { +// return &RequestContext{} +// }) +// if err != nil { +// // handle error +// } +// +// ctx, err := pool.Get() +// if err != nil { +// // handle error +// } +// ... +// pool.Put(ctx) // // Note: This is intentionally minimal and inlined for performance-sensitive paths. // If Go ever adds generics to sync.Pool (e.g. Go 1.23+), this becomes obsolete. import ( "errors" + "reflect" "sync" ) @@ -47,29 +52,58 @@ func NewLitePool[T any](newFn func() T) (*Pool[T], error) { } // Validate early that the result is non-nil test := newFn() - if any(test) == nil { + if isNilValue(test) { return nil, errors.New("litepool: constructor returned nil") } return &Pool[T]{ pool: sync.Pool{ New: func() any { - v := newFn() - if any(v) == nil { - // This should never happen after validation above, - // but keep the panic here as a last resort safety check - panic("litepool: constructor returned nil during runtime") - } - return v + return newFn() }, }, new: newFn, }, nil } -func (p *Pool[T]) Get() T { - //nolint:forcetypeassert // safe due to validated New - return p.pool.Get().(T) +func (p *Pool[T]) Get() (T, error) { + raw := p.pool.Get() + // sync.Pool.Get() returns nil when New is nil and the pool is empty. + // Guard here before the type assertion so interface-typed pools (Pool[any]) + // never panic on an untyped nil returned by a stateful factory. + if raw == nil { + var zero T + return zero, errors.New("litepool: factory returned nil") + } + v, ok := raw.(T) + if !ok { + var zero T + return zero, errors.New("litepool: pool returned unexpected type") + } + // A stateful factory can violate its contract after construction by returning + // a typed nil (e.g. (*T)(nil)). Return an error rather than panicking so + // callers can propagate it cleanly without crashing in-flight requests. + if isNilValue(v) { + var zero T + return zero, errors.New("litepool: factory returned nil") + } + return v, nil +} + +// isNilValue reports whether v is nil, handling both untyped nil (interface{} == nil) +// and typed nils (e.g. (*T)(nil) stored in an interface). The reflect path is only +// reached for nilable kinds; value types like structs always return false. +func isNilValue[T any](v T) bool { + if any(v) == nil { + return true + } + rv := reflect.ValueOf(any(v)) + switch rv.Kind() { + case reflect.Ptr, reflect.Interface, reflect.Map, reflect.Slice, reflect.Chan, reflect.Func: + return rv.IsNil() + default: + return false + } } func (p *Pool[T]) Put(v T) { diff --git a/pkg/pool/lite_pool_test.go b/pkg/pool/lite_pool_test.go new file mode 100644 index 00000000..201f9647 --- /dev/null +++ b/pkg/pool/lite_pool_test.go @@ -0,0 +1,115 @@ +package pool + +import ( + "sync" + "testing" +) + +func TestNewLitePool_NilConstructorReturnsError(t *testing.T) { + t.Parallel() + + _, err := NewLitePool[*int](nil) + if err == nil { + t.Fatal("expected error for nil constructor, got nil") + } +} + +func TestNewLitePool_GetPutRoundTrip(t *testing.T) { + t.Parallel() + + p, err := NewLitePool(func() *int { + v := 42 + return &v + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + got, err := p.Get() + if err != nil { + t.Fatalf("Get returned unexpected error: %v", err) + } + if got == nil { + t.Fatal("Get returned nil") + } + if *got != 42 { + t.Errorf("Get returned %d, want 42", *got) + } + p.Put(got) +} + +// TestLitePool_IsNilValue verifies that isNilValue correctly identifies typed and untyped nils. +func TestLitePool_IsNilValue(t *testing.T) { + t.Parallel() + + // A (*int)(nil) stored in an interface{} is not == nil at the interface level. + var typedNil *int + if !isNilValue(typedNil) { + t.Fatal("isNilValue must return true for a typed nil pointer") + } + + // Non-nil pointer must not be detected as nil. + v := 1 + if isNilValue(&v) { + t.Fatal("isNilValue must return false for a non-nil pointer") + } +} + +// TestLitePool_Get_ErrorOnNilFactory verifies that Get returns a non-nil error (not a panic) +// when a stateful factory violates its contract and returns nil after construction. +// We cannot replicate this via NewLitePool (it validates at construction time), so we +// construct a Pool directly with a sync.Pool whose New always yields a typed nil. +func TestLitePool_Get_ErrorOnNilFactory(t *testing.T) { + t.Parallel() + + // Bypass NewLitePool's construction-time guard so we can exercise the Get() error path. + p := &Pool[*int]{ + pool: sync.Pool{ + New: func() any { + // Return a typed nil (*int)(nil) - not the same as untyped nil in an + // interface, which is what isNilValue is specifically designed to catch. + var n *int + return n + }, + }, + } + + got, err := p.Get() + if err == nil { + t.Fatal("expected error when factory returns nil, got nil error") + } + if got != nil { + t.Errorf("expected zero value on error, got non-nil: %v", got) + } + const wantMsg = "litepool: factory returned nil" + if err.Error() != wantMsg { + t.Errorf("unexpected error message: got %q, want %q", err.Error(), wantMsg) + } +} + +// TestLitePool_Get_InterfaceTyped_UntypedNil verifies that Get on a Pool[any] does not +// panic when sync.Pool returns an untyped nil (raw == nil). The forced assertion path +// used to execute before the nil guard, so this test exercises the new ordering. +func TestLitePool_Get_InterfaceTyped_UntypedNil(t *testing.T) { + t.Parallel() + + // Pool[any] with a New that returns untyped nil - simulates a stateful factory that + // has exhausted its resource. sync.Pool can also return nil when New is not set and + // the pool is empty, but setting New here makes the behaviour deterministic. + p := &Pool[any]{ + pool: sync.Pool{ + New: func() any { + return nil // untyped nil; raw.(T) would panic before the nil check + }, + }, + } + + // Must return an error, never panic. + got, err := p.Get() + if err == nil { + t.Fatal("expected error for untyped nil from interface-typed pool, got nil error") + } + if got != nil { + t.Errorf("expected nil zero value on error, got: %v", got) + } +} diff --git a/test/cmd/ollamock/streaming.go b/test/cmd/ollamock/streaming.go index 800ed14b..4b93cb0d 100644 --- a/test/cmd/ollamock/streaming.go +++ b/test/cmd/ollamock/streaming.go @@ -310,6 +310,29 @@ func (srv *mockServer) streamOpenAIChat(w http.ResponseWriter, r *http.Request, } data, _ := json.Marshal(final) writeSSEData(w, data) + _ = rc.Flush() + + // Terminal usage chunk in the include_usage shape: choices:[] with populated usage. + // Matches the non-stream path token numbers so streaming and non-streaming agree. + // This exercises the choices:[] usage-capture path in the translation layer and + // gives the translator a real output_tokens value instead of falling back to the + // chars/4 estimate. + usageChunk := map[string]any{ + "id": id, + "object": "chat.completion.chunk", + "created": now, + "model": model, + "choices": []map[string]any{}, + "usage": map[string]any{ + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + }, + } + usageData, _ := json.Marshal(usageChunk) + writeSSEData(w, usageData) + _ = rc.Flush() + _, _ = fmt.Fprint(w, "data: [DONE]\n\n") _ = rc.Flush() }