import * as assert from 'node:assert';

let invocationEvents = new Map();

export default {

tailStream(onsetEvent, env, ctx) {

const invocationId = onsetEvent.invocationId;

const baseTimestamp = onsetEvent.timestamp.getTime();

// Initialize event list for this invocation

const events = [

{

type: onsetEvent.event.type,

relativeTimeMs: 0,

info: onsetEvent.event.info?.type,

entrypoint: onsetEvent.event.entrypoint,

},

];

invocationEvents.set(invocationId, { baseTimestamp, events });

return (event) => {

const relativeTimeMs = event.timestamp.getTime() - baseTimestamp;

const eventRecord = {

type: event.event.type,

relativeTimeMs,

};

// Capture additional info for specific event types

if (event.event.type === 'return') {

eventRecord.returnInfo = event.event.info;

} else if (event.event.type === 'outcome') {

eventRecord.outcome = event.event.outcome;

eventRecord.cpuTime = event.event.cpuTime;

eventRecord.wallTime = event.event.wallTime;

} else if (event.event.type === 'attributes') {

eventRecord.attributes = event.event.info;

}

events.push(eventRecord);

};

},

export const test = {

async test() {

// Wait briefly for all tail events to be processed

await scheduler.wait(200);

// Find the StreamingService invocation (the JSRPC call)

let streamingServiceEvents = null;

for (const [invocationId, data] of invocationEvents.entries()) {

const onset = data.events[0];

if (onset.entrypoint === 'StreamingService' && onset.info === 'jsrpc') {

streamingServiceEvents = data.events;

break;

}

}

if (!streamingServiceEvents) {

console.log('Captured invocations:');

for (const [invocationId, data] of invocationEvents.entries()) {

console.log(` ${invocationId}: ${JSON.stringify(data.events[0])}`);

}

throw new Error(

'Could not find StreamingService JSRPC invocation events'

);

}

console.log(

'StreamingService events:',

JSON.stringify(streamingServiceEvents, null, 2)

);

// Find the key events

const onset = streamingServiceEvents.find((e) => e.type === 'onset');

const returnEvent = streamingServiceEvents.find((e) => e.type === 'return');

const outcome = streamingServiceEvents.find((e) => e.type === 'outcome');

// Validate all events are present

assert.ok(onset, 'Missing onset event');

assert.ok(outcome, 'Missing outcome event');

if (!returnEvent) {

console.log('WARNING: No return event found for JSRPC invocation!');

console.log(

'This may indicate setReturn() is not being called for JSRPC.'

);

throw new Error('Return event not captured for JSRPC invocation');

}

const timeToReturn = returnEvent.relativeTimeMs;

const timeToOutcome = outcome.relativeTimeMs;

const gap = timeToOutcome - timeToReturn;

console.log(`timeToReturn: ${timeToReturn}ms`);

console.log(`timeToOutcome: ${timeToOutcome}ms`);

console.log(`Gap between Return and Outcome: ${gap}ms`);

// Key validation: Return must happen BEFORE Outcome

if (timeToReturn >= timeToOutcome) {

throw new Error(

`Return event (${timeToReturn}ms) should happen before Outcome (${timeToOutcome}ms). ` +

`This indicates Return is fired when stream ends, not when handler returns.`

);

}

// Return should happen at least 200ms after onset (handler sleeps for 500ms)

if (timeToReturn < 200) {

throw new Error(

`Return event (${timeToReturn}ms) happened too quickly. ` +

`Expected at least 200ms from onset (handler sleeps for 500ms before returning).`

);

}

// Return should be significantly before Outcome (at least 200ms gap for stream drain)

if (gap < 200) {

throw new Error(

`Gap between Return (${timeToReturn}ms) and Outcome (${timeToOutcome}ms) is only ${gap}ms. ` +

`Expected at least 200ms for stream consumption.`

);

}

console.log(`PASS: Return event occurs ${gap}ms before Outcome`);

},

import { WorkerEntrypoint } from 'cloudflare:workers';

export class StreamingService extends WorkerEntrypoint {

async getStreamWithDelays() {

// Sleep 500ms before returning the stream

await scheduler.wait(500);

// Return a ReadableStream that takes ~450ms to drain

const encoder = new TextEncoder();

let chunksSent = 0;

const stream = new ReadableStream({

async pull(controller) {

if (chunksSent >= 3) {

controller.close();

return;

}

// Sleep 150ms between chunks (3 chunks = ~450ms to drain)

await scheduler.wait(150);

controller.enqueue(encoder.encode(`chunk${chunksSent++}\n`));

},

});

return stream;

}

export default {

async test(controller, env, ctx) {

// Call the streaming RPC method

const stream = await env.StreamingService.getStreamWithDelays();

// Consume the stream fully

const reader = stream.getReader();

let chunks = [];

while (true) {

const { done, value } = await reader.read();

if (done) break;

chunks.push(new TextDecoder().decode(value));

}

// Verify we got all chunks

if (chunks.length !== 3) {

throw new Error(`Expected 3 chunks, got ${chunks.length}`);

}

// The actual timing validation is done in the tail worker's test() handler

},

using Workerd = import "/workerd/workerd.capnp";

# Test to validate timing semantics for JSRPC invocations with streaming responses.

# This test verifies that the Return event is emitted when the handler returns,

# NOT when the stream is fully consumed.

#

# Expected timeline:

# T=0: Onset event (invocation starts)

# T=~500ms: Return event (handler returns ReadableStream after 500ms sleep)

# T=~950ms: Outcome event (stream fully consumed after ~450ms more)

#

# If Return event occurs at T=~950ms instead of T=~500ms, this indicates a bug

# where setReturn() is called when capabilities are released rather than

# when the handler actually returns.

const unitTests :Workerd.Config = (

services = [

(name = "jsrpc-timing-test", worker = .mainWorker),

(name = "tail", worker = .tailWorker),

],

);

const mainWorker :Workerd.Worker = (

modules = [

(name = "worker", esModule = embed "jsrpc-timing-test.js")

],

compatibilityDate = "2024-01-01",

compatibilityFlags = ["nodejs_compat", "experimental", "precise_timers"],

bindings = [

(name = "StreamingService", service = (name = "jsrpc-timing-test", entrypoint = "StreamingService")),

],

streamingTails = ["tail"],

);

const tailWorker :Workerd.Worker = (

modules = [

(name = "worker", esModule = embed "jsrpc-timing-test-tail.js")

],

compatibilityDate = "2024-01-01",

compatibilityFlags = ["experimental", "nodejs_compat", "precise_timers"],

);

ioCtx(ioCtx),

// Override call() to emit the Return event when the top-level RPC call completes.

// This marks when the handler returned a value, NOT when all data has been streamed or all

// capabilities released.

kj::Promise<void> call(CallContext callContext) override {

return JsRpcTargetBase::call(kj::mv(callContext)).then([this]() {

KJ_IF_SOME(t, ioCtx.getWorkerTracer()) {

t.setReturn(ioCtx.now());

}

});

}

IoContext& ioCtx;