From 9b0d98132c65a3f62cfb0beb3e41aa5b2192af4e Mon Sep 17 00:00:00 2001
From: Selena Yang <179177246+selenayang888@users.noreply.github.com>
Date: Wed, 2 Sep 2026 18:28:05 -0700
Subject: [PATCH] fix bug to use model context length

---
 .../generative/chat/search_options.cc         | 12 ++++++---
 .../generative/chat/search_options.h          |  4 +--
 .../internal_api/chat/search_options_test.cc  | 25 +++++++++++++++++++
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc
index 4ae6df6a8..e472d14c1 100644
--- a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc
+++ b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc
@@ -17,17 +17,21 @@ int ApplySearchOptions(const SearchOptions& options,
                        ExecutionProvider ep,
                        bool use_full_context,
                        int default_max_output_tokens) {
-  // Determine model's max context length from genai_config.json search.max_length
+  // model.context_length is the model's input + output capacity. Older packages may only provide search.max_length.
   int model_max_length = 0;
-  if (config.search.has_value()) {
+  if (config.model.has_value()) {
+    model_max_length = config.model->context_length;
+  }
+  if (model_max_length <= 0 && config.search.has_value()) {
     model_max_length = config.search->max_length;
   }
 
   if (model_max_length <= 0) {
-    FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "model genai_config.json is missing search.max_length");
+    FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL,
+             "model genai_config.json is missing model.context_length and search.max_length");
   }
 
-  // genai_config.json's search.max_length (read above) is the source of truth for the total input+output budget.
+  // The resolved model context length above is the source of truth for the total input + output budget.
   // The catalog's maxOutputTokens is informational metadata only and is intentionally NOT used to clamp generation:
   // it is commonly a conservative 2048 that would wrongly cap larger contexts (e.g. the 3072 vision default). A
   // user-supplied max_output_tokens is honored as-is and only rejected if input+output exceeds max_length below.
diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h
index 7bc162c8a..27e17b298 100644
--- a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h
+++ b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h
@@ -50,13 +50,13 @@ struct SearchOptions {
 };
 
 /// Apply search options to OgaGeneratorParams.
-/// Validates token budget (input + output vs model max_length from config).
+/// Validates the input + output token budget against model.context_length, falling back to search.max_length.
 /// Returns the computed max_length that was set on the params.
 /// Throws fl::Exception on invalid configuration (e.g., input too long for model).
 ///
 /// @param options            Search options extracted from the request
 /// @param input_token_count  Number of tokens in the encoded prompt
-/// @param config             Model's GenAI config (for search.max_length)
+/// @param config             Model's GenAI config (for model.context_length and legacy search.max_length)
 /// @param gen_params         ORT GenAI generator params to configure
 /// @param ep                 Resolved execution provider. Used to enable chunked prefill by default
 ///                           on providers that benefit from it (CUDA, NvTensorRtRtx, WebGPU, CPU).
diff --git a/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc b/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc
index eb6ee5a32..bda1267c9 100644
--- a/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc
+++ b/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc
@@ -109,6 +109,31 @@ TEST_F(SearchOptionsTest, TokenBudgetExceededThrows) {
                fl::Exception);
 }
 
+TEST_F(SearchOptionsTest, ModelContextLengthTakesPrecedenceOverSearchMaxLength) {
+  SearchOptions opts;
+  opts.max_output_tokens = 256;
+  auto params = MakeParams();
+  GenAIConfig config;
+  config.model.emplace().context_length = 131072;
+  config.search.emplace().max_length = 4096;
+
+  int max_length = ApplySearchOptions(opts, 127826, config, *params, ExecutionProvider::kDefault);
+
+  EXPECT_EQ(max_length, 128082);
+}
+
+TEST_F(SearchOptionsTest, SearchMaxLengthSupportsLegacyConfig) {
+  SearchOptions opts;
+  opts.max_output_tokens = 256;
+  auto params = MakeParams();
+  GenAIConfig config;
+  config.search.emplace().max_length = 4096;
+
+  int max_length = ApplySearchOptions(opts, 3840, config, *params, ExecutionProvider::kDefault);
+
+  EXPECT_EQ(max_length, 4096);
+}
+
 TEST_F(SearchOptionsTest, TemperatureZeroDisablesSampling) {
   SearchOptions opts;
   opts.temperature = 0.0f;