fix: context window management

py/genai_client/message_builders/semoss_base/document_input.py

@@ -38,6 +38,10 @@ class DocumentInputError(ValueError):
     """An attachment could not be delivered without losing its contents."""
 
 
+class ContextBudgetError(DocumentInputError):
+    """The request is too large for the model's context; callers may compact and retry."""
+
+
 @dataclass(frozen=True)
 class ExtractedDocument:
     text: str
@@ -207,7 +211,7 @@ def prepare(
                     message.type = SEMOSSMessageType.INPUT_TEXT
 
             if extracted_chars > self.max_text_chars:
-                raise DocumentInputError(
+                raise ContextBudgetError(
                     "The attached documents contain too much text for one conversation "
                     f"({self.max_text_chars:,} character limit). Use fewer pages or files."
                 )
@@ -472,7 +476,7 @@ def _check_context_budget(
                 ]
             text_bytes += len(json.dumps(payload, ensure_ascii=False).encode("utf-8"))
         if text_bytes + int(output_tokens) > settings.context_window:
-            raise DocumentInputError(
+            raise ContextBudgetError(
                 "The extracted documents and conversation exceed the conservative "
                 "text budget for this model. Attach fewer pages, start a new conversation, "
                 "or reduce the requested output length. No content was truncated."

py/genai_client/text_generation/model_engine_exception.py

@@ -1,10 +1,25 @@
-from typing import Any
+from typing import Any, Optional
 import traceback, re, json
 import openai
 from pydantic import BaseModel
 from anthropic import APIStatusError, APIConnectionError, APITimeoutError
 from botocore.exceptions import ClientError, BotoCoreError
 
+from ..message_builders.semoss_base.document_input import ContextBudgetError
+
+# reason Java checks to compact/fork and retry instead of failing the room
+CONTEXT_OVERFLOW = "CONTEXT_OVERFLOW"
+
+# provider wording for "request exceeds the context window" (lowercase)
+_CONTEXT_OVERFLOW_PATTERNS = (
+    "maximum context length",  # openai / vllm / openai-compatible servers
+    "prompt is too long",  # anthropic
+    "input is too long",  # bedrock
+    "too many input tokens",  # bedrock
+    "exceeds the maximum number of tokens",  # google
+    "input token count",  # google
+)
+
 
 class AnthropicRefusalError(RuntimeError):
     """Raised when Anthropic returns stop_reason='refusal'."""
@@ -18,6 +33,9 @@ class ErrorDetails(BaseModel):
     client: str
     model: str
     traceback: str
+    reason: Optional[str] = None
+    # which rule set reason, e.g. "code:context_length_exceeded" or "text:prompt is too long"
+    reason_detail: Optional[str] = None
 
 
 class ModelEngineException:
@@ -28,6 +46,14 @@ def __init__(self, error: Any, client: str, model: str):
         self.traceback = traceback.format_exc()
 
     def parse_error(self) -> ErrorDetails:
+        details = self._parse_by_client()
+        rule = self._context_overflow_rule(details)
+        if rule:
+            details.reason = CONTEXT_OVERFLOW
+            details.reason_detail = rule
+        return details
+
+    def _parse_by_client(self) -> ErrorDetails:
         if self.client == "anthropic":
             return self._parse_anthropic_error()
         elif self.client in ["google", "vertex", "gemini"]:
@@ -44,6 +70,23 @@ def parse_error(self) -> ErrorDetails:
             error_type="Internal Server Error",
         )
 
+    def _context_overflow_rule(self, details: ErrorDetails) -> Optional[str]:
+        # structured signals first, provider wording as the fallback
+        if isinstance(self.error, ContextBudgetError):
+            return "type:ContextBudgetError"
+        if getattr(self.error, "code", None) == "context_length_exceeded":
+            return "code:context_length_exceeded"  # openai / azure
+        if isinstance(self.error, APIStatusError):
+            body = self.error.body if isinstance(self.error.body, dict) else {}
+            inner = body.get("error")
+            if isinstance(inner, dict) and inner.get("type") == "request_too_large":
+                return "type:request_too_large"  # anthropic 413, byte limit
+        text = f"{details.error_type} {details.message}".lower()
+        for pattern in _CONTEXT_OVERFLOW_PATTERNS:
+            if pattern in text:
+                return f"text:{pattern}"
+        return None
+
     def _parse_google_error(self) -> ErrorDetails:
         """
         Handles Google GenAI errors which often look like:

src/prerna/engine/api/IModelEngine.java

@@ -163,6 +163,14 @@ default MultiModalEmbeddingsModelEngineResponse multiModalEmbeddings(List<String
 	@IgnoreEngineLogging
 	int getContextWindow();
 
+	/**
+	 * @return configured max output tokens, or 0 when unknown
+	 */
+	@IgnoreEngineLogging
+	default long getMaxTokens() {
+		return 0;
+	}
+
 	// ------------------------------------------------------------------
 	// Batch model calls (native provider Batch APIs).
 	//

src/prerna/engine/impl/model/AbstractModelEngine.java

@@ -903,4 +903,9 @@ public boolean holdsFileLocks() {
 	public int getContextWindow() {
 		return this.contextWindow;
 	}
+
+	@Override
+	public long getMaxTokens() {
+		return this.maxTokens == null ? 0 : this.maxTokens;
+	}
 }

src/prerna/engine/impl/model/Room.java

@@ -613,6 +613,99 @@ public ResponseMessage appendHarnessResponse(String text, String parentMessageId
 		}
 	}
 
+	/**
+	 * Flags a message so provider requests replace tool payloads at and above it
+	 * with a placeholder. Stored history is unchanged. Returns false when the
+	 * message is missing or already flagged.
+	 */
+	public boolean markPruneToolsAbove(String messageId, Insight insight) {
+		ReentrantLock lock = getMessageLock();
+		lock.lock();
+		try {
+			String userId = insight.getUser().getPrimaryLoginToken().getId();
+			try (RoomMessageStore.RoomMutationLock ignored = RoomMessageStore.acquireMutationLock(this)) {
+				this.insight = insight;
+				RoomMessageStore.refreshFromLatestProjection(this, userId);
+				for (AbstractMessage message : messages) {
+					if (messageId.equals(message.getMessageId())) {
+						if (message.getPruneToolsAbove()) {
+							return false;
+						}
+						message.setPruneToolsAbove(true);
+						RoomMessageStore.persist(this, userId);
+						return true;
+					}
+				}
+				return false;
+			}
+		} finally {
+			lock.unlock();
+		}
+	}
+
+	/**
+	 * Adds a sibling of the newest tool-results message under the given tool call,
+	 * with the largest results replaced by a short note until about tokensToFree
+	 * (chars / 4) is freed; at least the largest result is replaced. The original
+	 * stays stored but off the active branch. Returns how many results were
+	 * replaced; 0 means nothing was forked.
+	 */
+	public int forkToolResultsWithStubs(String toolCallMessageId, long tokensToFree, int minStubChars,
+			Insight insight) {
+		ReentrantLock lock = getMessageLock();
+		lock.lock();
+		try {
+			String userId = insight.getUser().getPrimaryLoginToken().getId();
+			try (RoomMessageStore.RoomMutationLock ignored = RoomMessageStore.acquireMutationLock(this)) {
+				this.insight = insight;
+				RoomMessageStore.refreshFromLatestProjection(this, userId);
+				ToolExecutionContext context = findToolExecutionContext(
+						resolveToolContinuationMessageId(toolCallMessageId));
+				InputMessage original = findToolResultsMessage(context.toolResponse, context.toolResponseIdx);
+				if (original == null) {
+					return 0;
+				}
+				InputMessage fork = MessageUtils.deepCopy(original);
+				List<ToolResultPart> results = new ArrayList<>();
+				for (MessagePart part : fork.getParts()) {
+					if (part instanceof ToolResultMessagePart trPart && trPart.getToolResult() != null
+							&& trPart.getToolResult().getOutput() != null
+							&& trPart.getToolResult().getOutput().length() >= minStubChars) {
+						results.add(trPart.getToolResult());
+					}
+				}
+				results.sort((a, b) -> Integer.compare(b.getOutput().length(), a.getOutput().length()));
+				int replaced = 0;
+				long freed = 0;
+				for (ToolResultPart result : results) {
+					if (replaced > 0 && freed >= tokensToFree) {
+						break;
+					}
+					int chars = result.getOutput().length();
+					result.setOutput("[tool result too large for the model context: " + chars
+							+ " chars, not sent. Re-run with a narrower query or read it in parts"
+							+ " (e.g. ReadFile offset/limit).]");
+					freed += chars / 4;
+					replaced++;
+				}
+				if (replaced == 0) {
+					return 0;
+				}
+				fork.setMessageId(GUID.v7().toUUID().toString());
+				fork.setParentMessageId(original.getParentMessageId());
+				fork.setTokensInMessage(0);
+				fork.setRoom(this);
+				messages.add(fork);
+				RoomMessageStore.persist(this, userId);
+				classLogger.info("Forked tool results room={} original={} fork={} replaced={} freedTokens~{}", getId(),
+						original.getMessageId(), fork.getMessageId(), replaced, freed);
+				return replaced;
+			}
+		} finally {
+			lock.unlock();
+		}
+	}
+
 	/**
 	 * Persist a harness-owned completion after all tools finish, without another
 	 * model request.

src/prerna/engine/impl/model/message/MessageUtils.java

@@ -740,9 +740,30 @@ public static List<AbstractMessage> getMessageBranchFromParent(List<AbstractMess
 		return history;
 	}
 
+	/** Deep copy via the persistence serializer; keeps the room links and attachment bytes it skips. */
+	public static <T extends AbstractMessage> T deepCopy(T m) {
+		@SuppressWarnings("unchecked")
+		T copy = (T) GSON_FOR_DB.fromJson(GSON_FOR_DB.toJson(m), m.getClass());
+		copy.roomId = m.roomId;
+		copy.roomFolderPath = m.roomFolderPath;
+		// base64Data and the media room folder are not serialized; without them the
+		// provider gets attachments with no data
+		List<MessagePart> from = m.getParts();
+		List<MessagePart> to = copy.getParts();
+		if (from != null && to != null && from.size() == to.size()) {
+			for (int i = 0; i < from.size(); i++) {
+				if (from.get(i) instanceof MediaMessagePart src && to.get(i) instanceof MediaMessagePart dst
+						&& src.getMediaInfo() != null && dst.getMediaInfo() != null) {
+					dst.getMediaInfo().setBase64Data(src.getMediaInfo().getBase64Data());
+				}
+			}
+		}
+		return copy;
+	}
+
+	// only messages with tool parts change, so the rest stay as-is
 	private static AbstractMessage deepCopyForPruning(AbstractMessage m) {
-		String json = GSON_FOR_DB.toJson(m);
-		return GSON_FOR_DB.fromJson(json, m.getClass());
+		return m.hasToolResultPart() || m.hasToolCallPart() ? deepCopy(m) : m;
 	}
 
 	private static void pruneToolsPartFromMessage(AbstractMessage m) {

src/prerna/engine/impl/model/responses/AskErrorModelEngineResponse.java

@@ -36,12 +36,18 @@ public class AskErrorModelEngineResponse extends AskModelEngineResponse<String>
     public static final String CLIENT = "client";
     public static final String MODEL = "model";
     public static final String TRACEBACK = "traceback";
+    public static final String REASON = "reason";
+    public static final String REASON_DETAIL = "reason_detail";
+    // request did not fit the model context; callers may compact/fork and retry
+    public static final String REASON_CONTEXT_OVERFLOW = "CONTEXT_OVERFLOW";
 
     protected String errorType;
     protected int code;
     protected String client;
     protected String model;
     protected String traceback;
+    protected String reason;
+    protected String reasonDetail;
 
     public AskErrorModelEngineResponse(String message, String errorType, int code, String client, String model, String traceback) {
         super(message, 0, 0);
@@ -66,6 +72,16 @@ public String getStringResponse() {
 
     public String getTraceback() { return this.traceback; }
 
+    public String getReason() { return this.reason; }
+
+    public void setReason(String reason) { this.reason = reason; }
+
+    public String getReasonDetail() { return this.reasonDetail; }
+
+    public void setReasonDetail(String reasonDetail) { this.reasonDetail = reasonDetail; }
+
+    public boolean isContextOverflow() { return REASON_CONTEXT_OVERFLOW.equals(this.reason); }
+
     @Override
     public Map<String, Object> toMap() {
         Map<String, Object> map = super.toMap();
@@ -74,6 +90,8 @@ public Map<String, Object> toMap() {
         map.put(CLIENT, this.client);
         map.put(MODEL, this.model);
         map.put(TRACEBACK, this.traceback);
+        map.put(REASON, this.reason);
+        map.put(REASON_DETAIL, this.reasonDetail);
         return map;
     }
 }

src/prerna/engine/impl/model/responses/AskModelEngineResponse.java

@@ -231,8 +231,10 @@ public static AskModelEngineResponse<?> fromMap(Object responseObject) {
 			String model = safeString(modelResponse.get("model"));
 			String traceback = safeString(modelResponse.get("traceback"));
 
-			AskModelEngineResponse<?> errorResponse = new AskErrorModelEngineResponse(message, errorType, code, client,
-					model, traceback);
+			AskErrorModelEngineResponse errorResponse = new AskErrorModelEngineResponse(message, errorType, code,
+					client, model, traceback);
+			errorResponse.setReason(safeString(modelResponse.get(AskErrorModelEngineResponse.REASON)));
+			errorResponse.setReasonDetail(safeString(modelResponse.get(AskErrorModelEngineResponse.REASON_DETAIL)));
 			if (schemaVersion != null) {
 				errorResponse.setSchemaVersion(schemaVersion);
 			}

src/prerna/reactor/agent/runtime/HarnessToolExecutor.java

@@ -77,6 +77,8 @@
 import prerna.sablecc2.om.GenRowStruct;
 import prerna.sablecc2.om.PixelDataType;
 import prerna.sablecc2.om.ReactorKeysEnum;
+import prerna.engine.impl.model.responses.AskErrorModelEngineResponse;
+import prerna.sablecc2.om.execptions.SemossModelEngineException;
 import prerna.sablecc2.om.nounmeta.NounMetadata;
 
 /**
@@ -194,15 +196,32 @@ static ResponseMessage executeToolBatch(ResponseMessage toolResponse, AgentLoopS
                 return (ResponseMessage) room.getMessages().getLast();
             }
         }
+        // tool results are stored; prune older tool payloads if the context is near full
+        SemossAgentHarness.pruneToolsBeforeContinuation(ctx, parentMsgId);
         Map<String, Object> nextParams = new HashMap<>(paramMap);
         if (ctx.getAgentConfig().getFinishingTurns() > 0 && state.getIterations() >= ctx.getMaxTurns()) {
             nextParams.put("tool_choice", "none");
         }
         AgentRunStreamService.get().beginModelCall(ctx.getRunId());
         state.progress().beginModel();
         try {
-            nextModelResp = room.continueAfterToolExecutionResultsWithRuntimeContext(nextParams, parentMsgId,
-                    ctx.getModelEngine(), ctx.getInsight(), state.systemPrompt(), state.runtimeContext());
+            try {
+                nextModelResp = room.continueAfterToolExecutionResultsWithRuntimeContext(nextParams, parentMsgId,
+                        ctx.getModelEngine(), ctx.getInsight(), state.systemPrompt(), state.runtimeContext());
+            } catch (RuntimeException e) {
+                // too big for the context: fork with the largest results stubbed and retry once
+                AskErrorModelEngineResponse overflow = SemossModelEngineException.contextOverflowError(e);
+                if (overflow == null) {
+                    throw e;
+                }
+                logger.warn("HarnessToolExecutor: context overflow after tool batch room={} client={} rule={}",
+                        room.getId(), overflow.getClient(), overflow.getReasonDetail());
+                if (SemossAgentHarness.forkOversizedToolResults(ctx, parentMsgId) == 0) {
+                    throw e;
+                }
+                nextModelResp = room.continueAfterToolExecutionResultsWithRuntimeContext(nextParams, parentMsgId,
+                        ctx.getModelEngine(), ctx.getInsight(), state.systemPrompt(), state.runtimeContext());
+            }
         } finally {
             state.progress().endModel();
         }