Forward idempotency_key on AI assist so retries are not charged twice

dotnet/EcencyApi.Tests/AiAssistHandlerTests.cs

@@ -0,0 +1,80 @@
+using System.Text.Json.Nodes;
+using EcencyApi.Handlers;
+using EcencyApi.Infrastructure;
+using Xunit;
+
+namespace EcencyApi.Tests;
+
+/// <summary>
+/// AI assist bills whoever <c>us</c> names. ePoints returns the cached result for a
+/// repeated <c>idempotency_key</c> instead of charging again, so the proxy has to
+/// forward the key the client sent. Older clients omit it; those requests still go
+/// through.
+/// </summary>
+[Collection("ai-assist")]
+public class AiAssistHandlerTests : IDisposable
+{
+    private readonly CurationDeskTestSupport.Recorder _upstream = new();
+
+    public AiAssistHandlerTests()
+    {
+        // `code` "as:alice" validates as alice; anything else is invalid.
+        // Same extraction as production ValidateCode: a lone-surrogate escape is a
+        // real string, and GetValue<string>() throws on it.
+        PrivateApi.AiAssistValidateCode = body =>
+        {
+            string? code = null;
+            if (body["code"] is JsonValue codeValue && JsVal.TryGetStringLenient(codeValue, out var codeStr))
+                code = codeStr;
+            return Task.FromResult(code != null && code.StartsWith("as:", StringComparison.Ordinal) ? code[3..] : null);
+        };
+        PrivateApi.AiAssistUpstream = (endpoint, method, payload, _) =>
+            _upstream.Handle(endpoint, method, Array.Empty<KeyValuePair<string, string>>(), payload);
+    }
+
+    public void Dispose()
+    {
+        PrivateApi.AiAssistValidateCode = PrivateApi.ValidateCode;
+        PrivateApi.AiAssistUpstream = (endpoint, method, payload, timeoutMs) =>
+            EcencyApi.Infrastructure.ApiClient.ApiRequest(endpoint, method, null, payload, null, timeoutMs);
+    }
+
+    [Fact]
+    public async Task ForwardsTheIdempotencyKeyOnTheUpstreamBody()
+    {
+        var ctx = CurationDeskTestSupport.Post("/private-api/ai-assist", """
+            {"code":"as:alice","action":"summarize","text":"hello world","idempotency_key":"abcd1234efgh"}
+            """);
+        await PrivateApi.AiAssist(ctx);
+
+        Assert.Equal(200, ctx.Response.StatusCode);
+        var call = Assert.Single(_upstream.Calls);
+        Assert.Equal("ai-assist", call.Endpoint);
+        Assert.Equal(HttpMethod.Post, call.Method);
+        var payload = Assert.IsType<JsonObject>(call.Payload);
+        Assert.Equal("alice", payload["us"]?.GetValue<string>());
+        Assert.Equal("summarize", payload["action"]?.GetValue<string>());
+        Assert.Equal("hello world", payload["text"]?.GetValue<string>());
+        Assert.Equal("abcd1234efgh", payload["idempotency_key"]?.GetValue<string>());
+    }
+
+    [Fact]
+    public async Task ForwardsWhenTheClientOmitsTheIdempotencyKey()
+    {
+        var ctx = CurationDeskTestSupport.Post("/private-api/ai-assist", """
+            {"code":"as:alice","action":"summarize","text":"hello world"}
+            """);
+        await PrivateApi.AiAssist(ctx);
+
+        Assert.Equal(200, ctx.Response.StatusCode);
+        var call = Assert.Single(_upstream.Calls);
+        var payload = Assert.IsType<JsonObject>(call.Payload);
+        Assert.Equal("alice", payload["us"]?.GetValue<string>());
+        Assert.Equal("summarize", payload["action"]?.GetValue<string>());
+        Assert.Equal("hello world", payload["text"]?.GetValue<string>());
+        Assert.False(payload.ContainsKey("idempotency_key"));
+    }
+}
+
+[CollectionDefinition("ai-assist", DisableParallelization = true)]
+public class AiAssistCollection { }

dotnet/EcencyApi/Handlers/PrivateApi.Misc.cs

@@ -597,10 +597,21 @@ public static async Task AiAssistPrice(HttpContext ctx)
         await Upstream.Pipe(ApiClient.ApiRequest($"ai-assist-price?us={username}", HttpMethod.Get), ctx);
     }
 
+    /// <summary>Signed-code validation for AI assist, replaceable for tests (no chain RPC).</summary>
+    internal static Func<JsonObject, Task<string?>> AiAssistValidateCode = ValidateCode;
+
+    /// <summary>
+    /// The AI assist upstream call, replaceable so tests can observe the payload
+    /// without a network. The timeout is the long generation budget the handler passes.
+    /// </summary>
+    internal static Func<string, HttpMethod, JsonNode?, int, Task<UpstreamResponse>> AiAssistUpstream =
+        (endpoint, method, payload, timeoutMs) =>
+            ApiClient.ApiRequest(endpoint, method, null, payload, null, timeoutMs);
+
     public static async Task AiAssist(HttpContext ctx)
     {
         var body = await ctx.ReadBody();
-        var username = await ValidateCode(body);
+        var username = await AiAssistValidateCode(body);
         if (username == null)
         {
             await ctx.SendText(401, "Unauthorized");
@@ -610,9 +621,12 @@ public static async Task AiAssist(HttpContext ctx)
         {
             ["us"] = username,
         };
-        MiscCopyIfPresent(data, body, "action", "text");
+        // idempotency_key lets a retry recover the same paid assist instead of
+        // charging a second one. Older clients omit it, so it is copied only when
+        // present; the upstream validates its format itself.
+        MiscCopyIfPresent(data, body, "action", "text", "idempotency_key");
         // AI assist generation can take a long time; keep it long.
-        await Upstream.Pipe(ApiClient.ApiRequest("ai-assist", HttpMethod.Post, null, data, null, 120000), ctx);
+        await Upstream.Pipe(AiAssistUpstream("ai-assist", HttpMethod.Post, data, 120000), ctx);
     }
 
     public static async Task AiTranscribePrice(HttpContext ctx)

dotnet/parity/driver.py

@@ -349,6 +349,17 @@ def norm_body(text):
 # better reason: balances and APR move every block, so the run-vs-run comparison puts
 # it in `loose`, which still checks status and content-type. Add an entry here only
 # for a divergence that is deterministic and not already loose.
+#
+# Also not listed: /private-api/ai-assist::{min,pop,badcode}. Those catalog bodies
+# carry an empty or invalid `code`, so both this build and the reference image
+# answer 401 before any upstream call and the responses still match. A keyed retry
+# is different only after a valid code: this build copies `idempotency_key` onto
+# the upstream body when the client sent one (so a repeat recovers the same paid
+# assist instead of being charged again) and omits it when the client did not. The
+# reference image dropped the field either way. The catalog never sends a valid
+# code, so that difference is covered by AiAssistHandlerTests rather than an entry
+# here — an entry would skip status, content-type, and body and hide unrelated
+# regressions on the route.
 
 
 def diff(a_name, b_name, loose_name=None):