fix(ds4): enforce generation boundaries (#11821)

Clamp requested generation to the usable context after prompt sync while preserving the legacy 256-token fallback for omitted limits.

Constrain each speculative MTP cycle to the remaining request budget so accepted tokens cannot advance beyond the visible output limit.

Assisted-by: Codex:gpt-5.6-sol

Signed-off-by: Claudio Maradonna <git@codeshifter.xyz>

backend/cpp/ds4/generation_limits.h

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: MIT
+#pragma once
+
+#include <algorithm>
+
+namespace ds4cpp {
+
+inline int EffectiveGenerationLimit(int requested, int context_size,
+                                    int session_position) {
+  const int limit = requested > 0 ? requested : 256;
+  const int room = context_size - session_position;
+  if (room <= 1) return 0;
+  return std::min(limit, room - 1);
+}
+
+inline int RemainingGenerationBudget(int effective_limit, int produced) {
+  if (effective_limit <= produced) return 0;
+  return effective_limit - produced;
+}
+
+inline int SpeculativeAcceptedCapacity(int remaining, int draft_allowance,
+                                       int buffer_capacity) {
+  if (remaining <= 0 || draft_allowance < 0 || buffer_capacity <= 0) return 0;
+  return std::min({remaining, draft_allowance + 1, buffer_capacity});
+}
+
+} // namespace ds4cpp

backend/cpp/ds4/generation_limits_test.cpp

@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: MIT
+
+#include "generation_limits.h"
+
+#include <cstdio>
+
+namespace {
+
+int failures = 0;
+
+void check_equal(int got, int want, const char *name) {
+  if (got == want) return;
+  std::fprintf(stderr, "FAIL %s: got %d, want %d\n", name, got, want);
+  failures++;
+}
+
+// Mutation caught: treating omitted or negative max_tokens as unlimited instead
+// of preserving DS4's legacy 256-token default.
+void test_nonpositive_uses_legacy_default_when_space_permits() {
+  check_equal(ds4cpp::EffectiveGenerationLimit(0, 4096, 100), 256,
+              "zero max_tokens uses legacy default");
+  check_equal(ds4cpp::EffectiveGenerationLimit(-1, 4096, 100), 256,
+              "negative max_tokens uses legacy default");
+}
+
+// Mutation caught: applying the legacy default without clamping it to the
+// post-prefill context room and reserved slot.
+void test_legacy_default_is_clamped_by_context() {
+  check_equal(ds4cpp::EffectiveGenerationLimit(0, 300, 100), 199,
+              "legacy default is context-clamped");
+}
+
+// Mutation caught: allowing an explicitly large request to overrun the
+// post-prefill context boundary.
+void test_large_positive_limit_is_clamped_to_context() {
+  check_equal(ds4cpp::EffectiveGenerationLimit(32768, 32768, 100), 32667,
+              "large positive is context-clamped");
+}
+
+// Mutation caught: replacing every positive request with the legacy default
+// rather than preserving a smaller configured limit.
+void test_smaller_positive_limit_is_preserved() {
+  check_equal(ds4cpp::EffectiveGenerationLimit(64, 4096, 100), 64,
+              "smaller positive is preserved");
+}
+
+// Mutation caught: consuming the final context slot instead of reserving it as
+// required by DS4's generation loop.
+void test_no_usable_room_returns_zero() {
+  check_equal(ds4cpp::EffectiveGenerationLimit(32, 100, 99), 0,
+              "one remaining context slot is not usable");
+}
+
+// Mutation caught: sending the original generation limit to a later
+// speculative cycle instead of subtracting tokens already produced.
+void test_remaining_budget_accounts_for_produced_tokens() {
+  check_equal(ds4cpp::RemainingGenerationBudget(10, 4), 6,
+              "remaining budget subtracts produced tokens");
+  check_equal(ds4cpp::RemainingGenerationBudget(10, 12), 0,
+              "remaining budget never becomes negative");
+}
+
+// Mutation caught: giving speculative evaluation capacity beyond either the
+// output budget, the draft allowance plus its first target token, or the fixed
+// accepted-token buffer.
+void test_speculative_capacity_obeys_all_bounds() {
+  check_equal(ds4cpp::SpeculativeAcceptedCapacity(3, 8, 8), 3,
+              "capacity respects remaining output budget");
+  check_equal(ds4cpp::SpeculativeAcceptedCapacity(20, 4, 8), 5,
+              "capacity includes one target token beyond draft allowance");
+  check_equal(ds4cpp::SpeculativeAcceptedCapacity(20, 8, 6), 6,
+              "capacity respects fixed buffer");
+}
+
+} // namespace
+
+int main() {
+  test_nonpositive_uses_legacy_default_when_space_permits();
+  test_legacy_default_is_clamped_by_context();
+  test_large_positive_limit_is_clamped_to_context();
+  test_smaller_positive_limit_is_preserved();
+  test_no_usable_room_returns_zero();
+  test_remaining_budget_accounts_for_produced_tokens();
+  test_speculative_capacity_obeys_all_bounds();
+
+  if (failures == 0) {
+    std::fprintf(stderr, "all generation limit checks passed\n");
+    return 0;
+  }
+  std::fprintf(stderr, "%d check(s) failed\n", failures);
+  return 1;
+}

backend/cpp/ds4/grpc-server.cpp

@@ -10,6 +10,7 @@
 
 #include "dsml_parser.h"   // populated in Task 12
 #include "dsml_renderer.h" // populated in Task 16
+#include "generation_limits.h"
 #include "kv_cache.h"      // populated in Task 17
 
 extern "C" {
@@ -769,7 +770,6 @@ public:
         }
         ds4_tokens prompt = {};
         build_prompt(g_engine, request, &prompt);
-        int n_predict = request->tokens() > 0 ? request->tokens() : 256;
 
         const bool think_enabled = ds4_think_mode_enabled(parse_think_mode(request));
         const bool starts_in_thinking = think_enabled &&
@@ -792,6 +792,9 @@ public:
         int prompt_len = prompt.len;
         ds4_tokens_free(&prompt);
         if (rc == 0) {
+            const int n_predict = ds4cpp::EffectiveGenerationLimit(
+                request->tokens(), ds4_session_ctx(g_session),
+                ds4_session_pos(g_session));
             const int eos = ds4_token_eos(g_engine);
             const int draft_max = ds4_engine_mtp_draft_tokens(g_engine);
             int produced = 0;
@@ -810,9 +813,12 @@ public:
                 if (draft_max > 0 && sp.temperature <= 0.0f) {
                     constexpr int kAcceptedMax = 8;
                     int accepted[kAcceptedMax];
-                    int cap = std::min(kAcceptedMax, draft_max + 1);
+                    const int remaining = ds4cpp::RemainingGenerationBudget(
+                        n_predict, produced);
+                    const int cap = ds4cpp::SpeculativeAcceptedCapacity(
+                        remaining, draft_max, kAcceptedMax);
                     int n = ds4_session_eval_speculative_argmax(
-                        g_session, first, draft_max, eos,
+                        g_session, first, remaining, eos,
                         accepted, cap, err, sizeof(err));
                     if (n < 0) { rc = -1; break; }
                     bool stop = false;
@@ -873,7 +879,6 @@ public:
         }
         ds4_tokens prompt = {};
         build_prompt(g_engine, request, &prompt);
-        int n_predict = request->tokens() > 0 ? request->tokens() : 256;
 
         const bool think_enabled = ds4_think_mode_enabled(parse_think_mode(request));
         const bool starts_in_thinking = think_enabled &&
@@ -891,6 +896,9 @@ public:
         int rc = ds4_session_sync(g_session, &prompt, err, sizeof(err));
         ds4_tokens_free(&prompt);
         if (rc == 0) {
+            const int n_predict = ds4cpp::EffectiveGenerationLimit(
+                request->tokens(), ds4_session_ctx(g_session),
+                ds4_session_pos(g_session));
             const int eos = ds4_token_eos(g_engine);
             const int draft_max = ds4_engine_mtp_draft_tokens(g_engine);
             int produced = 0;
@@ -908,9 +916,12 @@ public:
                 if (draft_max > 0 && sp.temperature <= 0.0f) {
                     constexpr int kAcceptedMax = 8;
                     int accepted[kAcceptedMax];
-                    int cap = std::min(kAcceptedMax, draft_max + 1);
+                    const int remaining = ds4cpp::RemainingGenerationBudget(
+                        n_predict, produced);
+                    const int cap = ds4cpp::SpeculativeAcceptedCapacity(
+                        remaining, draft_max, kAcceptedMax);
                     int n = ds4_session_eval_speculative_argmax(
-                        g_session, first, draft_max, eos,
+                        g_session, first, remaining, eos,
                         accepted, cap, err, sizeof(err));
                     if (n < 0) { rc = -1; break; }
                     bool stop = false;

docs/content/advanced/model-configuration.md

@@ -173,6 +173,14 @@ These settings will be used as defaults for all the API calls to the model.
 | `tfz` | float | `1.0` | Tail free z parameter |
 | `keep` | int | `0` | Number of tokens to keep from the prompt |
 
+{{% notice note %}}
+The DS4 backend preserves its legacy behavior for omitted or non-positive
+`max_tokens` values by generating at most 256 tokens. Set `max_tokens` to a
+positive value when you need a specific DS4 output limit. After processing the
+prompt, DS4 clamps that limit to the available context space and reserves one
+context slot for safe generation.
+{{% /notice %}}
+
 ### Language and Translation
 
 | Field | Type | Description |