From 2adbd54014a57709854e6889c2e52d136907e946 Mon Sep 17 00:00:00 2001
From: NotXf1le <89696340+NotXf1le@users.noreply.github.com>
Date: Wed, 23 Sep 2026 14:32:31 +0200
Subject: [PATCH] Handle sampled OpenRouter label logprobs

---
 src/openrouter.ts         | 19 +++++++++++++++++++
 tests/openrouter.test.mjs | 27 ++++++++++++++++++++++++---
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/src/openrouter.ts b/src/openrouter.ts
index 96a8f44..6f2b9bf 100644
--- a/src/openrouter.ts
+++ b/src/openrouter.ts
@@ -132,6 +132,25 @@ function parseScores(value: unknown, candidates: readonly string[]): {
     found.set(entry.token, entry.logprob);
   }
 
+  const sampled = content[0];
+  if (typeof sampled.token === "string" && expected.has(sampled.token)) {
+    if (typeof sampled.logprob !== "number" || !Number.isFinite(sampled.logprob)
+      || sampled.logprob > 0 || sampled.logprob <= CLAMPED_LOGPROB) {
+      throw new ScoringError(`OpenRouter returned an invalid or clamped logprob for label ${sampled.token}.`);
+    }
+    if (sampled.bytes !== undefined && sampled.bytes !== null) {
+      if (!Array.isArray(sampled.bytes) || sampled.bytes.length !== 1
+        || sampled.bytes[0] !== sampled.token.charCodeAt(0)) {
+        throw new ScoringError(`OpenRouter returned invalid bytes for label ${sampled.token}.`);
+      }
+    }
+    const topLogprob = found.get(sampled.token);
+    if (topLogprob !== undefined && topLogprob !== sampled.logprob) {
+      throw new ScoringError(`OpenRouter returned conflicting logprobs for label ${sampled.token}.`);
+    }
+    found.set(sampled.token, sampled.logprob);
+  }
+
   if (found.size === 0) {
     throw new ScoringError("OpenRouter did not return logprobs for any choice label.");
   }
diff --git a/tests/openrouter.test.mjs b/tests/openrouter.test.mjs
index b904cef..f5f3b82 100644
--- a/tests/openrouter.test.mjs
+++ b/tests/openrouter.test.mjs
@@ -26,7 +26,9 @@ const unorderedTopLogprobs = Object.freeze([
 
 function scoredPosition(topLogprobs = unorderedTopLogprobs) {
   return {
-    token: "A", bytes: [65], logprob: -1.2338635921, top_logprobs: topLogprobs,
+    token: "A", bytes: [65],
+    logprob: topLogprobs.find(({ token }) => token === "A")?.logprob ?? -1.2338635921,
+    top_logprobs: topLogprobs,
   };
 }
 
@@ -64,7 +66,7 @@ function chooser(f, extra = {}) {
   });
 }
 
-test("uses first-position label logprobs instead of the sampled token", async () => {
+test("uses first-position label logprobs", async () => {
   const value = response();
   value.usage = {
     prompt_tokens: 180,
@@ -158,14 +160,33 @@ test("assigns zero probability to labels omitted from top logprobs", async () =>
     .reduce((sum, probability) => sum + probability, 0) - 1) < 1e-12);
 });
 
+test("includes the sampled label when it is absent from top logprobs", async () => {
+  const top = unorderedTopLogprobs.filter(({ token }) => token !== "C");
+  const value = response(top, {
+    logprobs: { content: [{ token: "C", bytes: [67], logprob: -1.1088635921, top_logprobs: top }] },
+  });
+  const decision = await chooser(fixture(value))(request);
+
+  assert.equal(decision.choice, "billing");
+  assert.equal(decision.scores.billing, -1.1088635921);
+  assert.ok(decision.distribution.billing > 0);
+});
+
 const malformed = [
   ["null logprobs", response(undefined, { logprobs: null }), /logprobs/i],
   ["no scored position", response(undefined, { logprobs: { content: [] } }), /exactly one/i],
   ["multiple scored positions", response(undefined, {
     logprobs: { content: [scoredPosition(), scoredPosition()] },
   }), /exactly one/i],
-  ["no choice labels", response([{ token: "x", bytes: [120], logprob: -0.1 }]),
+  ["no choice labels", response([{ token: "x", bytes: [120], logprob: -0.1 }], {
+    logprobs: { content: [{ token: "x", bytes: [120], logprob: -0.1,
+      top_logprobs: [{ token: "x", bytes: [120], logprob: -0.1 }] }] },
+  }),
     /any choice label/i],
+  ["conflicting sampled logprob", response(unorderedTopLogprobs, {
+    logprobs: { content: [{ token: "A", bytes: [65], logprob: -1.5,
+      top_logprobs: unorderedTopLogprobs }] },
+  }), /conflicting.*A|A.*conflicting/i],
   ["duplicate required label", response([...unorderedTopLogprobs, unorderedTopLogprobs[0]]), /duplicate.*A|A.*duplicate/i],
   ["invalid logprob", response(unorderedTopLogprobs.map((entry) =>
     entry.token === "C" ? { ...entry, logprob: 0.1 } : entry)), /logprob/i],