fix(agent): harden structured stream guardrails

Author: Bill Leoutsakos

apps/sim/executor/execution/block-executor.test.ts

@@ -6,6 +6,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { clearLargeValueCacheForTests } from '@/lib/execution/payloads/cache'
 import { isLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest-metadata'
 import { isLargeValueRef } from '@/lib/execution/payloads/large-value-ref'
+import { REDACTION_FAILED_MARKER, redactObjectStrings } from '@/lib/logs/execution/pii-redaction'
 import { projectTraceSpansForSecrets } from '@/lib/logs/execution/trace-secret-projection'
 import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
 import { BlockType, EDGE } from '@/executor/constants'
@@ -1565,6 +1566,79 @@ describe('BlockExecutor streaming pump', () => {
     )
   })
 
+  it('fails empty structured streams when the terminal event reports a token limit', async () => {
+    const onFullContent = vi.fn()
+    const handler = createAgentEventsStreamingHandler({
+      events: [{ type: 'turn_end', turn: 'final', finishReason: 'length' }],
+      onFullContent,
+    })
+    const { executor, block, state } = createExecutor(handler)
+    block.config.params = {
+      responseFormat: { type: 'object', properties: { answer: { type: 'string' } } },
+    }
+    const ctx = createContext(state)
+
+    await expect(executor.execute(ctx, createNode(block), block)).rejects.toThrow(
+      /maximum output-token limit/i
+    )
+
+    expect(onFullContent).not.toHaveBeenCalled()
+    expect(state.getBlockOutput(block.id)).toMatchObject({
+      content: '',
+      error: expect.stringMatching(/maximum output-token limit/i),
+    })
+    expect(state.getBlockOutput(block.id)?.providerTiming?.timeSegments?.[0]?.finishReason).toBe(
+      'length'
+    )
+  })
+
+  it('preserves the original failure when provider-diagnostic redaction scrubs', async () => {
+    const redactor = vi.mocked(redactObjectStrings)
+    redactor
+      .mockImplementationOnce(async (value, options) => {
+        expect(options.onFailure).toBe('throw')
+        return value as never
+      })
+      .mockImplementationOnce(async (value, options) => {
+        expect(options.onFailure).toBe('scrub')
+        return {
+          ...(value as Record<string, unknown>),
+          content: REDACTION_FAILED_MARKER,
+          model: REDACTION_FAILED_MARKER,
+        } as never
+      })
+
+    const onFullContent = vi.fn()
+    const handler = createAgentEventsStreamingHandler({
+      events: [{ type: 'text_delta', text: '{"answer":"unfinished', turn: 'final' }],
+      finishReason: 'max_tokens',
+      onFullContent,
+    })
+    const { executor, block, state } = createExecutor(handler)
+    block.config.params = {
+      responseFormat: { type: 'object', properties: { answer: { type: 'string' } } },
+    }
+    const ctx = createContext(state)
+    ctx.piiBlockOutputRedaction = {
+      enabled: true,
+      entityTypes: ['EMAIL_ADDRESS'],
+      language: 'en',
+    }
+
+    await expect(executor.execute(ctx, createNode(block), block)).rejects.toThrow(
+      /maximum output-token limit/i
+    )
+
+    expect(redactor).toHaveBeenCalledTimes(2)
+    expect(onFullContent).not.toHaveBeenCalled()
+    expect(state.getBlockOutput(block.id)).toMatchObject({
+      content: REDACTION_FAILED_MARKER,
+      model: REDACTION_FAILED_MARKER,
+      tokens: { input: 1, output: 2, total: 3 },
+      error: expect.stringMatching(/maximum output-token limit/i),
+    })
+  })
+
   it('soft-completes on user abort with drained answer text (no failed block)', async () => {
     const abortController = new AbortController()
     const handler = createAgentEventsStreamingHandler({

apps/sim/executor/execution/block-executor.ts

@@ -76,6 +76,7 @@ import {
   type VariableResolver,
 } from '@/executor/variables/resolver'
 import { createAgentStreamPump } from '@/providers/stream-pump'
+import { enrichLastModelSegment } from '@/providers/trace-enrichment'
 import type { SerializedBlock } from '@/serializer/types'
 import { SYSTEM_SUBBLOCK_IDS } from '@/triggers/constants'
 
@@ -637,7 +638,7 @@ export class BlockExecutor {
     let providerDiagnostics: Record<string, unknown> = {}
     for (const key of ['content', 'model', 'tokens', 'toolCalls', 'providerTiming', 'cost']) {
       const value = failureDiagnosticOutput?.[key]
-      if (value !== undefined && (key !== 'content' || value !== '')) {
+      if (value !== undefined) {
         providerDiagnostics[key] = value
       }
     }
@@ -647,7 +648,7 @@ export class BlockExecutor {
         entityTypes: ctx.piiBlockOutputRedaction.entityTypes,
         language: ctx.piiBlockOutputRedaction.language,
         customPatterns: ctx.piiBlockOutputRedaction.customPatterns,
-        onFailure: 'throw',
+        onFailure: 'scrub',
       })
     }
     Object.assign(errorOutput, providerDiagnostics)
@@ -1242,11 +1243,8 @@ export class BlockExecutor {
     }
 
     let fullContent = pumpResult.answerText
-    if (!fullContent) {
-      return
-    }
 
-    if (piiEnabled && ctx.piiBlockOutputRedaction) {
+    if (fullContent && piiEnabled && ctx.piiBlockOutputRedaction) {
       // Mask before writing to `execution.output` or `onFullContent`.
       fullContent = await redactObjectStrings(fullContent, {
         entityTypes: ctx.piiBlockOutputRedaction.entityTypes,
@@ -1257,15 +1255,28 @@ export class BlockExecutor {
     }
 
     const executionOutput = streamingExec.execution?.output
+    if (pumpResult.finishReason && executionOutput?.providerTiming?.timeSegments) {
+      enrichLastModelSegment(executionOutput.providerTiming.timeSegments, {
+        finishReason: pumpResult.finishReason,
+      })
+    }
+    if (executionOutput && typeof executionOutput === 'object' && parsedResponseFormat) {
+      // Retain even empty content for failed-block diagnostics, but reject it
+      // before parsing so token-limited structured data never reaches downstream.
+      executionOutput.content = fullContent
+      assertStructuredOutputNotTokenLimited(executionOutput.providerTiming, executionOutput)
+    }
+
+    if (!fullContent) {
+      return
+    }
+
     if (executionOutput && typeof executionOutput === 'object') {
       let parsedForFormat = false
       if (responseFormat) {
         // Retain the drained text for failed-block diagnostics, but reject it
         // before parsing so truncated structured data never reaches downstream.
         executionOutput.content = fullContent
-        if (parsedResponseFormat) {
-          assertStructuredOutputNotTokenLimited(executionOutput.providerTiming, executionOutput)
-        }
         try {
           const parsed = JSON.parse(fullContent.trim())
           streamingExec.execution.output = {

apps/sim/providers/bedrock/utils.stream.test.ts

@@ -36,12 +36,18 @@ describe('createReadableStreamFromBedrockStream', () => {
         yield {
           metadata: { usage: { inputTokens: 2, outputTokens: 3 } },
         } as any
+        yield {
+          messageStop: { stopReason: 'max_tokens' },
+        } as any
       })(),
       onComplete
     )
 
     const events = await collectEvents(stream)
-    expect(events).toEqual([{ type: 'text_delta', text: 'Done', turn: 'final' }])
+    expect(events).toEqual([
+      { type: 'text_delta', text: 'Done', turn: 'final' },
+      { type: 'turn_end', turn: 'final', finishReason: 'max_tokens' },
+    ])
     expect(events.some((e) => e.type === 'thinking_delta')).toBe(false)
     expect(events.some((e) => e.type === 'tool_call_start')).toBe(false)
     expect(onComplete).toHaveBeenCalledWith('Done', { inputTokens: 2, outputTokens: 3 })

apps/sim/providers/bedrock/utils.ts

@@ -42,6 +42,7 @@ export function createReadableStreamFromBedrockStream(
   let fullContent = ''
   let inputTokens = 0
   let outputTokens = 0
+  let finishReason: string | undefined
   let cancelled = false
   let streamIterator: AsyncIterator<ConverseStreamOutput> | undefined
 
@@ -55,6 +56,9 @@ export function createReadableStreamFromBedrockStream(
           const event = next.value
           const streamError = getBedrockStreamError(event)
           if (streamError) throw streamError
+          if (event.messageStop?.stopReason) {
+            finishReason = event.messageStop.stopReason
+          }
           if (event.contentBlockDelta?.delta?.text) {
             const text = event.contentBlockDelta.delta.text
             fullContent += text
@@ -69,6 +73,9 @@ export function createReadableStreamFromBedrockStream(
         if (onComplete) {
           onComplete(fullContent, { inputTokens, outputTokens })
         }
+        if (finishReason) {
+          controller.enqueue({ type: 'turn_end', turn: 'final', finishReason })
+        }
 
         controller.close()
       } catch (err) {

apps/sim/providers/google/utils.stream.test.ts

@@ -26,6 +26,7 @@ describe('createReadableStreamFromGeminiStream', () => {
         yield {
           candidates: [
             {
+              finishReason: 'MAX_TOKENS',
               content: {
                 parts: [
                   { text: 'Reasoning step. ', thought: true },
@@ -48,6 +49,7 @@ describe('createReadableStreamFromGeminiStream', () => {
     expect(events).toEqual([
       { type: 'thinking_delta', text: 'Reasoning step. ' },
       { type: 'text_delta', text: 'Final answer.', turn: 'final' },
+      { type: 'turn_end', turn: 'final', finishReason: 'MAX_TOKENS' },
     ])
     expect(onComplete).toHaveBeenCalledWith(
       'Final answer.',

apps/sim/providers/google/utils.ts

@@ -252,6 +252,7 @@ export function createReadableStreamFromGeminiStream(
     cachedContentTokenCount: 0,
     totalTokenCount: 0,
   }
+  let finishReason: string | undefined
   let cancelled = false
   let streamIterator: AsyncIterator<GenerateContentResponse> | undefined
 
@@ -275,8 +276,12 @@ export function createReadableStreamFromGeminiStream(
           if (chunk.usageMetadata) {
             usage = convertUsageMetadata(chunk.usageMetadata)
           }
+          const candidate = chunk.candidates?.[0]
+          if (candidate?.finishReason) {
+            finishReason = String(candidate.finishReason)
+          }
 
-          const parts = chunk.candidates?.[0]?.content?.parts
+          const parts = candidate?.content?.parts
           if (Array.isArray(parts)) {
             for (const part of parts) {
               if (!part.text) continue
@@ -301,6 +306,9 @@ export function createReadableStreamFromGeminiStream(
 
         if (cancelled) return
         onComplete?.(fullContent, usage, fullThinking || undefined)
+        if (finishReason) {
+          controller.enqueue({ type: 'turn_end', turn: 'final', finishReason })
+        }
         controller.close()
       } catch (error) {
         if (!cancelled) {

apps/sim/providers/openai-compat/stream-events.test.ts

@@ -57,7 +57,9 @@ describe('createOpenAICompatibleAgentEventStream', () => {
       { providerName: 'Groq' }
     )
     const events = await collectEvents(stream)
-    expect(events.every((e) => e.type === 'text_delta')).toBe(true)
+    expect(events.filter((e) => e.type !== 'turn_end').every((e) => e.type === 'text_delta')).toBe(
+      true
+    )
     expect(events.some((e) => e.type === 'thinking_delta')).toBe(false)
   })
 
@@ -141,6 +143,9 @@ describe('createOpenAICompatibleAgentEventStream', () => {
     const stream = createOpenAICompatibleAgentEventStream(
       (async function* () {
         yield* openaiCompatTextOnlyChunks as any
+        yield {
+          choices: [{ delta: {}, finish_reason: 'length' }],
+        } as any
       })(),
       { providerName: 'DeepSeek', onComplete }
     )
@@ -152,6 +157,7 @@ describe('createOpenAICompatibleAgentEventStream', () => {
         .join('')
     ).toBe('Hello world')
     expect(onComplete.mock.calls[0][0].content).toBe('Hello world')
+    expect(events).toContainEqual({ type: 'turn_end', turn: 'final', finishReason: 'length' })
   })
 
   it('surfaces documented in-band provider errors', async () => {

apps/sim/providers/openai-compat/stream-events.ts

@@ -270,6 +270,9 @@ export function createOpenAICompatibleAgentEventStream(
             ...(finishReason ? { finishReason } : {}),
           })
         }
+        if (finishReason) {
+          controller.enqueue({ type: 'turn_end', turn, finishReason })
+        }
 
         controller.close()
       } catch (error) {

apps/sim/providers/openai/utils.stream.test.ts

@@ -126,7 +126,10 @@ describe('createReadableStreamFromResponses', () => {
 
     const events = await collectEvents(createReadableStreamFromResponses(response, onComplete))
 
-    expect(events).toEqual([{ type: 'text_delta', text: 'Truncated answer', turn: 'final' }])
+    expect(events).toEqual([
+      { type: 'text_delta', text: 'Truncated answer', turn: 'final' },
+      { type: 'turn_end', turn: 'final', finishReason: 'max_output_tokens' },
+    ])
     expect(onComplete).toHaveBeenCalledWith(
       'Truncated answer',
       {

apps/sim/providers/openai/utils.ts

@@ -438,6 +438,7 @@ export function createReadableStreamFromResponses(
         let fullThinking = ''
         let finalUsage: ResponsesUsageTokens | undefined
         let completed = false
+        let finishReason: string | undefined
         let sawFunctionCall = false
 
         try {
@@ -464,6 +465,7 @@ export function createReadableStreamFromResponses(
                 throw new Error(`OpenAI Responses stream incomplete: ${reason}`)
               }
               finalUsage = parseResponsesUsage(event.response.usage)
+              finishReason = reason
               completed = true
               continue
             }
@@ -499,6 +501,9 @@ export function createReadableStreamFromResponses(
           }
 
           onComplete?.(fullContent, finalUsage, fullThinking || undefined)
+          if (finishReason) {
+            controller.enqueue({ type: 'turn_end', turn: 'final', finishReason })
+          }
           controller.close()
         } catch (error) {
           if (!streamAbortController.signal.aborted) {