Skip to content

Feature Request: Access to StreamState from SSEAdapter #553

@zngly-vlad

Description

@zngly-vlad

What we are building

We have a custom adapter, SSEVercelAIAdapter, which extends VercelAIAdapter to implement a richer version of the Vercel AI SDK Data Stream Protocol. A core part of what it does is track the precise timing of each content phase — recording when a text block started being streamed and when it ended, and doing the same for reasoning blocks. These startTime / endTime values are emitted to the client as providerMetadata on the text-end and reasoning-end events.

The pain point

The adapter is the only place in the system that knows when each content block started and finished streaming. It computes this timing by tracking state across chunks as they flow through transform(). But this timing data lives and dies entirely within the adapter — it is emitted to the client and then discarded. It never makes it back to the content blocks in StreamState.

What we actually want is for startTime and endTime to be attached as metadata directly onto the TextContent and ReasoningContent blocks that StreamState is accumulating — so that when those blocks are serialised and saved to the database, the timing information is preserved alongside the content itself. Right now the content blocks have no metadata field at all, and even if they did, the adapter has no reference to StreamState and no ability to write back into the blocks being built by the provider.

Side note: Huge thanks for the work behind this. It’s thoughtful, polished, and very enjoyable to use 🙌

<?php

use NeuronAI\Chat\Messages\Stream\Adapters\VercelAIAdapter;
use NeuronAI\Chat\Messages\Stream\Chunks\ReasoningChunk;
use NeuronAI\Chat\Messages\Stream\Chunks\TextChunk;
use NeuronAI\Chat\Messages\Stream\Chunks\ToolCallChunk;
use NeuronAI\Chat\Messages\Stream\Chunks\ToolResultChunk;
use NeuronAI\UniqueIdGenerator;

/**
 * Adapter for Vercel AI SDK Data Stream Protocol.
 *
 * @see https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol
 */
class MY_CUSTOM_VercelAIAdapter extends VercelAIAdapter {
    protected bool $started = false;
    protected string $messageId = '';

    protected bool $textStarted = false;
    protected int $textStartTime = 0;
    protected string $textId = '';

    protected bool $reasoningStarted = false;
    protected int $reasoningStartTime = 0;
    protected string $reasoningId = '';

    /** @var array<string, string> */
    protected array $toolCallIds = [];

    public function transform(object $chunk): iterable {
        // Lazy init on the first chunk
        if (!$this->started) {
            $this->started = true;
            $this->messageId = $chunk->messageId ?? UniqueIdGenerator::generateId();
            yield $this->sse(['type' => 'start', 'messageId' => $this->messageId]);
        }

        yield from match (true) {
            $chunk instanceof TextChunk => $this->handleText($chunk),
            $chunk instanceof ReasoningChunk => $this->handleReasoning($chunk),
            $chunk instanceof ToolCallChunk => $this->handleToolCall($chunk),
            $chunk instanceof ToolResultChunk => $this->handleToolResult($chunk),
            default => []
        };
    }

    /**
     * @return iterable<string>
     */
    protected function handleReasoningStart(): iterable {
        if (!$this->reasoningStarted) {
            $this->reasoningStarted = true;
            $this->reasoningId = UniqueIdGenerator::generateId();
            $this->reasoningStartTime = time();
            yield $this->sse([
                'type' => 'reasoning-start',
                'id' => $this->reasoningId
            ]);
        }
    }

    /**
     * @return iterable<string>
     */
    protected function handleReasoning(ReasoningChunk $chunk): iterable {
        yield from $this->handleTextEnd();
        yield from $this->handleReasoningStart();

        yield $this->sse([
            'type' => 'reasoning-delta',
            'id' => $this->reasoningId,
            'delta' => $chunk->content,
        ]);
    }

    /**
     * @return iterable<string>
     */
    protected function handleReasoningEnd(): iterable {
        if ($this->reasoningStarted) {
            $this->reasoningStarted = false;
            yield $this->sse([
                'type' => 'reasoning-end',
                'id' => $this->reasoningId,
                'providerMetadata' => [
                    'data' => [
                        'startTime' => $this->reasoningStartTime,
                        'endTime' => time(),
                    ]
                ],
            ]);
            $this->reasoningId = '';
            $this->reasoningStartTime = 0;
        }
    }

    /**
     * @return iterable<string>
     */
    protected function handleTextStart(TextChunk $chunk): iterable {
        if (!$this->textStarted) {
            $this->textStarted = true;
            $this->textId = UniqueIdGenerator::generateId();
            $this->textStartTime = time();
            yield $this->sse([
                'type' => 'text-start',
                'id' => $this->textId,
            ]);
        }
    }

    /**
     * @return iterable<string>
     */
    protected function handleText(TextChunk $chunk): iterable {
        // some weird case when we are still reasonging and the text chunk content is an empty string
        // in that case we should not end the reasoning but just skip the chunk
        if ($this->reasoningStarted && trim($chunk->content) === '') {
            return [];
        }

        yield from $this->handleReasoningEnd();
        yield from $this->handleTextStart($chunk);

        yield $this->sse([
            'type' => 'text-delta',
            'id' => $this->textId,
            'delta' => $chunk->content,
        ]);
    }

    /**
     * @return iterable<string>
     */
    protected function handleTextEnd(): iterable {
        if ($this->textStarted) {
            $this->textStarted = false;
            yield $this->sse([
                'type' => 'text-end',
                'id' => $this->textId,
                'providerMetadata' => [
                    'data' => [
                        'startTime' => $this->textStartTime,
                        'endTime' => time(),
                    ],
                ],
            ]);
            $this->textId = '';
            $this->textStartTime = 0;
        }
    }

    public function end(): iterable {
        yield from $this->handleTextEnd();

        yield from parent::end();
    }
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions