Chunk Assembly

Contents

Overview

flowchart TD
  SSE["Raw SSE Event Stream"]
  OAI["openai-stream-wrappers.ts"]
  GGL["google-stream-wrappers.ts"]
  ANT["anthropic-stream-wrappers.ts"]
  BDR["bedrock-stream-wrappers.ts"]
  NORM["Normalized Chunk"]
  ROUTER["Content Router"]
  TEXT["Text Accumulator"]
  TOOL["Tool-Call Accumulator"]
  THINK["Thinking Stream"]
  USAGE["usage-accumulator.ts"]
  BUILD["payloads.ts — buildEmbeddedRunPayloads()"]
  OUT["EmbeddedRunPayload"]

  SSE --> OAI
  SSE --> GGL
  SSE --> ANT
  SSE --> BDR
  OAI --> NORM
  GGL --> NORM
  ANT --> NORM
  BDR --> NORM
  NORM --> ROUTER
  ROUTER -->|"delta.content"| TEXT
  ROUTER -->|"delta.tool_calls"| TOOL
  ROUTER -->|"delta.thinking"| THINK
  NORM --> USAGE
  TEXT --> BUILD
  TOOL --> BUILD
  THINK --> BUILD
  USAGE --> BUILD
  BUILD --> OUT

The diagram above shows the full chunk-assembly data flow. Each provider wrapper sits behind a common interface: it consumes raw bytes from the HTTP response body and emits normalized chunk objects. The content router inspects each chunk's type discriminator and fans out to the appropriate accumulator. Finally, buildEmbeddedRunPayloads() merges all accumulated state into a single payload object.

Decision Points

How It Works

Chunk assembly operates in three logical phases: normalization, routing, and payload construction. Each phase is designed to be stateless with respect to the others, communicating only through the normalized chunk interface.

Phase 1 — Provider Normalization

Each LLM provider streams responses in a vendor-specific SSE format. The provider wrappers translate these into a common NormalizedChunk shape. OpenAI sends data: {choices: [{delta: {content, tool_calls}}]}, while Anthropic uses content_block_delta events with a type discriminator. Google's Vertex AI uses yet another event structure, and Bedrock wraps everything in an AWS EventStream binary framing.

// openai-stream-wrappers.ts — normalize OpenAI delta
interface OpenAIDelta {
  choices: Array<{
    delta: {
      content?: string;
      tool_calls?: Array<{
        index: number;
        id?: string;
        function?: { name?: string; arguments?: string };
      }>;
    };
    finish_reason: string | null;
  }>;
  usage?: { prompt_tokens: number; completion_tokens: number };
}

function normalizeOpenAIChunk(raw: OpenAIDelta): NormalizedChunk {
  const choice = raw.choices[0];
  return {
    content: choice.delta.content ?? null,
    toolCalls: choice.delta.tool_calls?.map(tc => ({
      index: tc.index,
      id: tc.id ?? null,
      name: tc.function?.name ?? null,
      argumentFragment: tc.function?.arguments ?? "",
    })) ?? null,
    thinking: null,
    finishReason: choice.finish_reason,
    usage: raw.usage ?? null,
  };
}

Phase 2 — Content Routing

Once normalized, each chunk is routed by type. Text content is appended to a running string buffer. Tool-call fragments are dispatched to per-index accumulators that build up the argument JSON incrementally. Thinking content is piped to a separate stream handled by thinking.ts.

// Simplified content routing logic
function routeChunk(chunk: NormalizedChunk, state: AssemblyState): void {
  if (chunk.thinking) {
    state.thinkingStream.push(chunk.thinking);
    return;
  }

  if (chunk.content) {
    state.textBuffer += chunk.content;
  }

  if (chunk.toolCalls) {
    for (const tc of chunk.toolCalls) {
      if (!state.toolAccumulators[tc.index]) {
        state.toolAccumulators[tc.index] = {
          id: tc.id ?? "",
          name: tc.name ?? "",
          argumentChunks: [],
        };
      }
      const acc = state.toolAccumulators[tc.index];
      if (tc.id) acc.id = tc.id;
      if (tc.name) acc.name = tc.name;
      acc.argumentChunks.push(tc.argumentFragment);
    }
  }

  if (chunk.usage) {
    state.usageAccumulator.add(chunk.usage);
  }
}

Phase 3 — Tool-Call Argument Repair

Tool-call arguments arrive as JSON fragments split across multiple SSE chunks. Because the split point is arbitrary (it may land in the middle of a string value or between object keys), naive concatenation can produce invalid JSON. The attempt.tool-call-normalization.ts module first joins all fragments and attempts JSON.parse(). If that fails, attempt.tool-call-argument-repair.ts applies heuristic fixes: closing unclosed strings, adding missing brackets, and stripping trailing commas.

Phase 4 — Payload Construction

After the stream terminates (or a tool-call is deemed complete), payloads.ts calls buildEmbeddedRunPayloads() to merge the text buffer, finalized tool calls, thinking content, and accumulated usage into an EmbeddedRunPayload. This payload is the unit of work consumed by the attempt loop and ultimately returned to the caller.

Behavior Diagram

sequenceDiagram
  participant Provider as LLM Provider
  participant Wrapper as Stream Wrapper
  participant Router as Content Router
  participant TextAcc as Text Accumulator
  participant ToolAcc as Tool-Call Accumulator
  participant Repair as Argument Repair
  participant Builder as Payload Builder

  Provider->>Wrapper: SSE chunk (vendor format)
  Wrapper->>Router: NormalizedChunk
  alt text content
    Router->>TextAcc: append(delta.content)
  else tool-call fragment
    Router->>ToolAcc: accumulate(index, fragment)
  end
  Provider->>Wrapper: SSE chunk #2
  Wrapper->>Router: NormalizedChunk
  Router->>ToolAcc: accumulate(index, fragment)
  Note over ToolAcc: JSON still incomplete
  Provider->>Wrapper: data: [DONE]
  Wrapper->>Router: finishReason = "stop"
  ToolAcc->>Repair: tryParse(joinedArgs)
  alt parse succeeds
    Repair->>Builder: validToolCall
  else parse fails
    Repair->>Repair: heuristic fix + retry
    Repair->>Builder: repairedToolCall
  end
  TextAcc->>Builder: fullText
  Builder->>Builder: buildEmbeddedRunPayloads()
  Builder-->>Provider: EmbeddedRunPayload

The sequence diagram above traces a typical two-chunk interaction where the first SSE event carries text content and the start of a tool call, the second continues the tool-call arguments, and the [DONE] sentinel triggers final assembly. Notice how the argument-repair step only activates when JSON.parse() fails on the concatenated fragments.

Annotated Code Walkthrough

Walk through the key functions in the order they execute during a single streaming response.

// stream-controller.ts — entry point
async function* processStream(
  providerType: ProviderType,
  response: ReadableStream<Uint8Array>,
): AsyncGenerator<NormalizedChunk> {
  const wrapper = getStreamWrapper(providerType);
  const reader = response.getReader();
  const decoder = new TextDecoder();
  let buffer = "";

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    buffer += decoder.decode(value, { stream: true });
    const events = extractSSEEvents(buffer);
    buffer = events.remainder;
    for (const event of events.parsed) {
      yield wrapper.normalize(event);
    }
  }
}

// attempt.tool-call-argument-repair.ts
function repairToolCallArguments(raw: string): object | null {
  // Attempt 1: direct parse
  try { return JSON.parse(raw); } catch {}

  // Attempt 2: close unclosed strings and brackets
  let fixed = raw;
  const openBraces = (fixed.match(/\{/g) || []).length;
  const closeBraces = (fixed.match(/\}/g) || []).length;
  if (openBraces > closeBraces) {
    fixed += "}".repeat(openBraces - closeBraces);
  }
  // Strip trailing commas before closing brace
  fixed = fixed.replace(/,\s*\}/g, "}");
  try { return JSON.parse(fixed); } catch {}

  // Attempt 3: close unclosed string literals
  const quoteCount = (fixed.match(/(?

// payloads.ts — final assembly
function buildEmbeddedRunPayloads(state: AssemblyState): EmbeddedRunPayload {
  const toolCalls = Object.values(state.toolAccumulators).map(acc => ({
    id: acc.id,
    name: acc.name,
    arguments: repairToolCallArguments(acc.argumentChunks.join("")),
  }));

  return {
    assistantText: state.textBuffer,
    toolCalls: toolCalls.filter(tc => tc.arguments !== null),
    malformedToolCalls: toolCalls.filter(tc => tc.arguments === null),
    thinking: state.thinkingStream.join(""),
    usage: state.usageAccumulator.summary(),
  };
}

src/agents/pi-embedded-runner/
├── streaming/
│   ├── openai-stream-wrappers.ts      ← OpenAI SSE normalization
│   ├── google-stream-wrappers.ts      ← Google Vertex AI normalization
│   ├── anthropic-stream-wrappers.ts   ← Anthropic SSE normalization
│   ├── bedrock-stream-wrappers.ts     ← AWS Bedrock EventStream normalization
│   ├── payloads.ts                    ← buildEmbeddedRunPayloads() assembly
│   ├── thinking.ts                    ← Thinking/reasoning content stream
│   └── usage-accumulator.ts           ← Token usage tracking across chunks
├── attempt/
│   ├── attempt.tool-call-normalization.ts  ← Tool-call JSON join + parse
│   └── attempt.tool-call-argument-repair.ts ← Heuristic JSON repair
└── stream-controller.ts              ← SSE event loop and provider dispatch