Address CodeRabbit provider feedback

apps/workshop/homarr_provider.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"log"
 	"net/http"
 	"net/url"
 	"os"
@@ -33,6 +34,8 @@ const (
 
 var errProviderResponseTooLarge = errors.New("provider response is too large")
 
+var clientControlledRoutingFields = []string{"models", "provider", "route", "plugins", "transforms"}
+
 type homarrProvider struct {
 	apiKey     string
 	baseURL    string
@@ -180,8 +183,7 @@ func (provider *homarrProvider) chat(event *core.RequestEvent) error {
 	if err := json.Unmarshal(body, &payload); err != nil {
 		return event.JSON(http.StatusBadRequest, openAIError("The request body is not valid OpenAI chat JSON."))
 	}
-	payload["model"] = openRouterModelID
-	payload["usage"] = map[string]any{"include": true}
+	sanitizeProviderPayload(payload)
 	upstreamBody, err := json.Marshal(payload)
 	if err != nil {
 		return event.JSON(http.StatusBadRequest, openAIError("The request body could not be prepared."))
@@ -206,7 +208,7 @@ func (provider *homarrProvider) chat(event *core.RequestEvent) error {
 		bytes.NewReader(upstreamBody),
 	)
 	if err != nil {
-		provider.finishActivity(event.App, privateRequest, publicActivity, "failed", requestUsage{}, startedAt)
+		provider.finishActivityWithRefund(event.App, privateRequest, publicActivity, "failed", requestUsage{}, startedAt, requestUnits)
 		return event.JSON(http.StatusInternalServerError, openAIError("The Homarr provider request could not be prepared."))
 	}
 	upstreamRequest.Header.Set("Authorization", "Bearer "+provider.apiKey)
@@ -218,10 +220,10 @@ func (provider *homarrProvider) chat(event *core.RequestEvent) error {
 
 	upstreamResponse, err := provider.httpClient.Do(upstreamRequest)
 	if err != nil {
-		provider.finishActivity(event.App, privateRequest, publicActivity, "failed", requestUsage{}, startedAt)
+		provider.finishActivityWithRefund(event.App, privateRequest, publicActivity, "failed", requestUsage{}, startedAt, requestUnits)
 		return event.JSON(http.StatusBadGateway, openAIError("The model endpoint could not be reached."))
 	}
-	defer upstreamResponse.Body.Close()
+	defer func() { _ = upstreamResponse.Body.Close() }()
 
 	event.Response.Header().Set("Cache-Control", "no-store")
 	event.Response.Header().Set("X-Accel-Buffering", "no")
@@ -238,7 +240,8 @@ func (provider *homarrProvider) chat(event *core.RequestEvent) error {
 	}
 	if upstreamResponse.StatusCode < 200 || upstreamResponse.StatusCode >= 300 {
 		errorBody, _ := io.ReadAll(io.LimitReader(upstreamResponse.Body, maxUpstreamErrorBytes))
-		provider.finishActivity(event.App, privateRequest, publicActivity, "failed", requestUsage{}, startedAt)
+		provider.finishActivityWithRefund(event.App, privateRequest, publicActivity, "failed", requestUsage{}, startedAt, requestUnits)
+		event.Response.Header().Set("X-Homarr-Quota-Remaining", strconv.Itoa(min(snapshot.Limit, snapshot.Remaining+requestUnits)))
 		return event.Blob(upstreamResponse.StatusCode, contentType, errorBody)
 	}
 
@@ -261,6 +264,14 @@ func (provider *homarrProvider) chat(event *core.RequestEvent) error {
 	return event.Blob(upstreamResponse.StatusCode, contentType, responseBody)
 }
 
+func sanitizeProviderPayload(payload map[string]any) {
+	for _, field := range clientControlledRoutingFields {
+		delete(payload, field)
+	}
+	payload["model"] = openRouterModelID
+	payload["usage"] = map[string]any{"include": true}
+}
+
 func readBoundedBody(reader io.Reader, limit int64) ([]byte, error) {
 	body, err := io.ReadAll(io.LimitReader(reader, limit+1))
 	if err != nil {
@@ -382,9 +393,21 @@ func (provider *homarrProvider) finishActivity(
 	status string,
 	usage requestUsage,
 	startedAt time.Time,
+) {
+	provider.finishActivityWithRefund(app, privateRequest, publicActivity, status, usage, startedAt, 0)
+}
+
+func (provider *homarrProvider) finishActivityWithRefund(
+	app core.App,
+	privateRequest *core.Record,
+	publicActivity *core.Record,
+	status string,
+	usage requestUsage,
+	startedAt time.Time,
+	refundUnits int,
 ) {
 	durationMs := max(0, int(provider.now().UTC().Sub(startedAt).Milliseconds()))
-	_ = app.RunInTransaction(func(txApp core.App) error {
+	err := app.RunInTransaction(func(txApp core.App) error {
 		for _, record := range []*core.Record{privateRequest, publicActivity} {
 			record.Set("status", status)
 			record.Set("inputTokens", usage.InputTokens)
@@ -407,11 +430,17 @@ func (provider *homarrProvider) finishActivity(
 		if !requestBelongsToQuotaDay(quota.GetString("day"), startedAt) {
 			return nil
 		}
+		if refundUnits > 0 {
+			quota.Set("used", max(0, quota.GetInt("used")-refundUnits))
+		}
 		quota.Set("inputTokens", quota.GetInt("inputTokens")+usage.InputTokens)
 		quota.Set("outputTokens", quota.GetInt("outputTokens")+usage.OutputTokens)
 		quota.Set("totalTokens", quota.GetInt("totalTokens")+usage.TotalTokens)
 		return txApp.Save(quota)
 	})
+	if err != nil {
+		log.Printf("Homarr provider accounting update failed for status %q", status)
+	}
 }
 
 func requestBelongsToQuotaDay(day string, startedAt time.Time) bool {
@@ -462,7 +491,7 @@ func (capture *sseUsageCapture) Write(chunk []byte) (int, error) {
 		capture.pending = capture.pending[end+1:]
 	}
 	if len(capture.pending) > maxUpstreamErrorBytes {
-		capture.pending = capture.pending[len(capture.pending)-maxUpstreamErrorBytes:]
+		capture.pending = nil
 	}
 	return len(chunk), nil
 }

apps/workshop/homarr_provider_test.go

@@ -1,12 +1,38 @@
 package main
 
 import (
+	"encoding/json"
 	"math"
 	"strings"
 	"testing"
 	"time"
 )
 
+func TestSanitizeProviderPayload(t *testing.T) {
+	payload := map[string]any{
+		"model":      homarrProviderModelID,
+		"models":     []string{"attacker/model"},
+		"provider":   map[string]any{"order": []string{"attacker"}},
+		"route":      "fallback",
+		"plugins":    []string{"web"},
+		"transforms": []string{"middle-out"},
+		"messages":   []any{map[string]any{"role": "user", "content": "hello"}},
+	}
+	sanitizeProviderPayload(payload)
+	for _, field := range clientControlledRoutingFields {
+		if _, exists := payload[field]; exists {
+			t.Fatalf("client-controlled routing field %q was forwarded", field)
+		}
+	}
+	if payload["model"] != openRouterModelID {
+		t.Fatalf("unexpected upstream model: %v", payload["model"])
+	}
+	usage, err := json.Marshal(payload["usage"])
+	if err != nil || string(usage) != `{"include":true}` {
+		t.Fatalf("unexpected usage options: %s, %v", usage, err)
+	}
+}
+
 func TestCountRequestUnits(t *testing.T) {
 	tests := []struct {
 		name     string

apps/workshop/pb_migrations/1786500000_homarr_provider.js

@@ -50,10 +50,10 @@ migrate(
         {
           type: "relation",
           name: "publicActivity",
-          required: true,
+          required: false,
           maxSelect: 1,
           collectionId: activity.id,
-          cascadeDelete: true,
+          cascadeDelete: false,
         },
       ],
     });

apps/workshop/tests/openrouter-mock.mjs

@@ -28,12 +28,23 @@ const server = createServer(async (request, response) => {
     json(response, 400, { error: { message: "Unexpected upstream model" } });
     return;
   }
+  const forbiddenRoutingField = ["models", "provider", "route", "plugins", "transforms"].find((field) => field in body);
+  if (forbiddenRoutingField) {
+    json(response, 400, { error: { message: `Client routing field was forwarded: ${forbiddenRoutingField}` } });
+    return;
+  }
+  if (body.messages?.some((message) => message.content === "upstream failure")) {
+    json(response, 503, { error: { message: "Simulated upstream failure" } });
+    return;
+  }
 
   const completion = {
     id: "gen-workshop-test",
     object: "chat.completion",
     model: body.model,
-    choices: [{ index: 0, message: { role: "assistant", content: "Hello from the Homarr provider" }, finish_reason: "stop" }],
+    choices: [
+      { index: 0, message: { role: "assistant", content: "Hello from the Homarr provider" }, finish_reason: "stop" },
+    ],
     usage: { prompt_tokens: 12, completion_tokens: 7, total_tokens: 19, cost: 0.00041 },
   };
   if (!body.stream) {

apps/workshop/tests/workshop.integration.mjs

@@ -229,12 +229,13 @@ while (true) {
   if (chunk.done) break;
   streamText += new TextDecoder().decode(chunk.value);
 }
-const streamFinishedAt = performance.now();
+const dataFrames = streamText.split(/\r?\n/u).filter((line) => line.startsWith("data:"));
 if (
   firstChunk.done ||
   !streamText.includes("Hello from ") ||
   !streamText.includes("the Homarr provider") ||
-  streamFinishedAt - firstChunkAt < 50 ||
+  dataFrames.length < 3 ||
+  dataFrames.at(-1)?.trim() !== "data: [DONE]" ||
   firstChunkAt - streamStartedAt > 5_000
 ) {
   throw new Error("The Homarr provider did not preserve the upstream SSE stream");
@@ -260,6 +261,11 @@ const toolResponse = await fetch(`${baseUrl}/api/ai/v1/chat/completions`, {
       { role: "tool", tool_call_id: "two", content: "{}" },
     ],
     stream: false,
+    models: ["attacker/model"],
+    provider: { order: ["attacker"] },
+    route: "fallback",
+    plugins: ["web"],
+    transforms: ["middle-out"],
   }),
 });
 if (!toolResponse.ok || toolResponse.headers.get("x-homarr-quota-remaining") !== "46") {
@@ -271,6 +277,23 @@ if (afterToolsUsage.used !== 4 || afterToolsUsage.remaining !== 46) {
   throw new Error(`Tool request accounting is incorrect: ${JSON.stringify(afterToolsUsage)}`);
 }
 
+const upstreamFailure = await fetch(`${baseUrl}/api/ai/v1/chat/completions`, {
+  method: "POST",
+  headers: { "content-type": "application/json", ...authorSession.headers },
+  body: JSON.stringify({
+    model: "homarr/deepseek-v4-flash-latest",
+    messages: [{ role: "user", content: "upstream failure" }],
+    stream: false,
+  }),
+});
+if (upstreamFailure.status !== 503 || upstreamFailure.headers.get("x-homarr-quota-remaining") !== "46") {
+  throw new Error("Failed upstream requests must refund their Homarr provider allowance");
+}
+const afterFailureUsage = await request("/api/ai/usage", { headers: authorSession.headers });
+if (afterFailureUsage.used !== 4 || afterFailureUsage.remaining !== 46) {
+  throw new Error(`Failed upstream request was charged: ${JSON.stringify(afterFailureUsage)}`);
+}
+
 const publicActivities = await request("/api/collections/assistant_activity/records?sort=-created&perPage=10");
 if (
   publicActivities.items.length < 2 ||
@@ -306,6 +329,18 @@ if (
 ) {
   throw new Error("Private provider accounting must identify the user without storing request content");
 }
+const requestToRetain = privateRequests.items.find((item) => item.publicActivity);
+if (!requestToRetain) throw new Error("Expected a private request linked to public activity");
+await request(`/api/collections/assistant_activity/records/${requestToRetain.publicActivity}`, {
+  method: "DELETE",
+  headers: rootHeaders,
+});
+const retainedRequest = await request(`/api/collections/assistant_requests/records/${requestToRetain.id}`, {
+  headers: rootHeaders,
+});
+if (retainedRequest.id !== requestToRetain.id || retainedRequest.publicActivity) {
+  throw new Error("Pruning public activity must retain private accounting records");
+}
 const quotas = await request("/api/collections/assistant_quotas/records?perPage=20", { headers: rootHeaders });
 const authorQuota = quotas.items.find((item) => item.user === author.id);
 if (!authorQuota || authorQuota.dailyLimit !== 50 || authorQuota.used !== 4 || authorQuota.totalTokens !== 38) {