diff --git a/.changeset/calm-timers-observe.md b/.changeset/calm-timers-observe.md new file mode 100644 index 0000000000..18116f4bea --- /dev/null +++ b/.changeset/calm-timers-observe.md @@ -0,0 +1,5 @@ +--- +'@workflow/world-vercel': patch +--- + +Add first-write WebSocket phase timings to stream write spans. diff --git a/packages/world-vercel/src/http-core.ts b/packages/world-vercel/src/http-core.ts index 888d22ee1f..4b01f2a5a1 100644 --- a/packages/world-vercel/src/http-core.ts +++ b/packages/world-vercel/src/http-core.ts @@ -17,7 +17,7 @@ * which is safe: `http-client.ts` imports nothing from `utils.ts` but a type. */ -import type { Span } from '@opentelemetry/api'; +import type { Attributes, Span } from '@opentelemetry/api'; import { getVercelOidcToken } from '@vercel/oidc'; import { EntityConflictError, @@ -429,7 +429,7 @@ export interface HttpClientSpanOptions { */ spanName?: string; /** Extra attributes merged on top of the standard HTTP attributes. */ - attributes?: Record; + attributes?: Attributes; } /** diff --git a/packages/world-vercel/src/ws-stream-session.test.ts b/packages/world-vercel/src/ws-stream-session.test.ts index 42acc8cab3..00ae1a0d4d 100644 --- a/packages/world-vercel/src/ws-stream-session.test.ts +++ b/packages/world-vercel/src/ws-stream-session.test.ts @@ -6,6 +6,7 @@ const { getVercelOidcToken, injectTraceContextIntoHeaders, sockets, + writeSpans, } = vi.hoisted(() => { const sockets: FakeSocket[] = []; class FakeSocket { @@ -70,6 +71,7 @@ const { getVercelOidcToken: vi.fn().mockResolvedValue(undefined), injectTraceContextIntoHeaders: vi.fn(), sockets, + writeSpans: [] as Array>, }; }); @@ -79,6 +81,24 @@ vi.mock('./telemetry.js', async (importOriginal) => { const actual = await importOriginal(); return { ...actual, injectTraceContextIntoHeaders }; }); +vi.mock('./http-core.js', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + withHttpClientSpan: vi.fn(async (options, callback) => { + if (options.spanName !== 'workflow.stream.write') { + return callback(undefined); + } + const attributes = { ...options.attributes }; + writeSpans.push(attributes); + return callback({ + setAttributes(next: Record) { + Object.assign(attributes, next); + }, + }); + }), + }; +}); const { createStreamWriteSession } = await import('./ws-stream-session.js'); @@ -99,6 +119,7 @@ beforeEach(() => { sockets.length = 0; getVercelOidcToken.mockClear(); injectTraceContextIntoHeaders.mockClear(); + writeSpans.length = 0; delete process.env.WORKFLOW_STREAMS_TRANSPORT; }); @@ -133,6 +154,51 @@ describe('v1 stream WebSocket writer lifecycle', () => { expect(closeHttp).toHaveBeenCalledTimes(1); }); + it('tags the first session and connection write with phase timings', async () => { + process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws'; + const { session } = makeSession(); + const writing = session.write(0, ['one']); + await vi.waitFor(() => expect(sockets).toHaveLength(1)); + sockets[0].open(); + await vi.waitFor(() => expect(sockets[0].sent).toHaveLength(1)); + sockets[0].reply( + encodeFrame({ type: 'write_ack', reqId: 1 }, new Uint8Array()) + ); + await writing; + + expect(writeSpans).toHaveLength(1); + expect(writeSpans[0]).toMatchObject({ + 'workflow.stream.ws.session_first_write': true, + 'workflow.stream.ws.connection_first_write': true, + 'workflow.stream.ws.connection_attempt': 1, + }); + for (const attribute of [ + 'workflow.stream.ws.session_to_write_ms', + 'workflow.stream.ws.write_wait_for_open_ms', + 'workflow.stream.ws.write_to_send_ms', + 'workflow.stream.ws.open_to_send_ms', + 'workflow.stream.ws.connect_ms', + 'workflow.stream.ws.config_token_ms', + 'workflow.stream.ws.send_to_reply_ms', + 'workflow.stream.ws.reply_processing_ms', + 'workflow.stream.ws.write_total_ms', + ]) { + expect(writeSpans[0][attribute]).toEqual(expect.any(Number)); + expect(writeSpans[0][attribute]).toBeGreaterThanOrEqual(0); + } + + const second = session.write(1, ['two']); + await vi.waitFor(() => expect(sockets[0].sent).toHaveLength(2)); + sockets[0].reply( + encodeFrame({ type: 'write_ack', reqId: 2 }, new Uint8Array()) + ); + await second; + expect(writeSpans[1]).toEqual({ + 'workflow.stream.transport': 'ws', + 'workflow.stream.ws.req_id': 2, + }); + }); + it('carries the first write when the socket opens inside the budget', async () => { process.env.WORKFLOW_STREAMS_TRANSPORT = 'ws'; const { session, writeHttp } = makeSession(); @@ -350,6 +416,12 @@ describe('v1 stream WebSocket writer lifecycle', () => { expect(sockets).toHaveLength(1); expect(writeHttp).not.toHaveBeenCalled(); expect(closeHttp).not.toHaveBeenCalled(); + expect(writeSpans[0]).toMatchObject({ + 'workflow.stream.ws.session_first_write': true, + 'workflow.stream.ws.send_to_reply_ms': expect.any(Number), + 'workflow.stream.ws.reply_processing_ms': expect.any(Number), + 'workflow.stream.ws.write_total_ms': expect.any(Number), + }); }); it('drains admitted work before reconnecting queued writes', async () => { @@ -393,6 +465,20 @@ describe('v1 stream WebSocket writer lifecycle', () => { await second; expect(writeHttp).not.toHaveBeenCalled(); + expect(writeSpans).toHaveLength(2); + expect(writeSpans[0]).toMatchObject({ + 'workflow.stream.ws.session_first_write': true, + 'workflow.stream.ws.connection_first_write': true, + 'workflow.stream.ws.connection_attempt': 1, + }); + expect(writeSpans[1]).toMatchObject({ + 'workflow.stream.ws.session_first_write': false, + 'workflow.stream.ws.connection_first_write': true, + 'workflow.stream.ws.connection_attempt': 2, + }); + expect(writeSpans[1]['workflow.stream.ws.connect_ms']).toEqual( + expect.any(Number) + ); }); it('requests fresh auth after an auth-expiry drain', async () => { diff --git a/packages/world-vercel/src/ws-stream-session.ts b/packages/world-vercel/src/ws-stream-session.ts index 0cdf388ed4..0c7ea7b2cd 100644 --- a/packages/world-vercel/src/ws-stream-session.ts +++ b/packages/world-vercel/src/ws-stream-session.ts @@ -1,3 +1,4 @@ +import type { Attributes, Span } from '@opentelemetry/api'; import { getVercelOidcToken } from '@vercel/oidc'; import type { StreamWriteSession } from '@workflow/world'; import type { WebSocket } from 'ws'; @@ -26,7 +27,117 @@ import { import { isWsStreamsTransportEnabled } from './ws-transport-enabled.js'; type Mode = 'connecting' | 'draining' | 'ws' | 'http' | 'closed' | 'poisoned'; +type WriteTiming = { + startedAt: number; + sessionFirstWrite: boolean; +}; +type ConnectionTiming = { + attempt: number; + configStartedAt: number; + configFinishedAt?: number; + socketStartedAt?: number; + openedAt?: number; + firstWriteSent: boolean; +}; +type PendingRequest = { + reqId: number; + resolve: (meta: Record) => void; + reject: (error: unknown) => void; + timer: ReturnType; + sentAt?: number; + replyReceivedAt?: number; +}; const MAX_IDLE_RECONNECTS = 3; + +function now(): number { + return performance.now(); +} + +function firstWriteAttributes( + write: WriteTiming | undefined, + connection: ConnectionTiming | undefined, + sessionCreatedAt: number +): Attributes { + if (!write) return {}; + return { + 'workflow.stream.ws.session_first_write': write.sessionFirstWrite, + 'workflow.stream.ws.connection_first_write': + connection?.firstWriteSent === false, + 'workflow.stream.ws.connection_attempt': connection?.attempt ?? 0, + 'workflow.stream.ws.session_to_write_ms': + write.startedAt - sessionCreatedAt, + }; +} + +function recordFirstWriteSetup( + span: Span | undefined, + write: WriteTiming, + connection: ConnectionTiming | undefined +): void { + if (connection) connection.firstWriteSent = true; + const openedAt = connection?.openedAt; + const socketStartedAt = connection?.socketStartedAt; + const configFinishedAt = connection?.configFinishedAt; + span?.setAttributes({ + 'workflow.stream.ws.write_wait_for_open_ms': openedAt + ? Math.max(0, openedAt - write.startedAt) + : 0, + ...(socketStartedAt && configFinishedAt + ? { + 'workflow.stream.ws.connect_setup_ms': + socketStartedAt - configFinishedAt, + } + : {}), + ...(socketStartedAt && openedAt + ? { 'workflow.stream.ws.connect_ms': openedAt - socketStartedAt } + : {}), + ...(configFinishedAt && connection + ? { + 'workflow.stream.ws.config_token_ms': + configFinishedAt - connection.configStartedAt, + } + : {}), + }); +} + +function recordFirstWriteSend( + span: Span | undefined, + write: WriteTiming, + connection: ConnectionTiming | undefined, + pending: PendingRequest +): void { + const sentAt = now(); + pending.sentAt = sentAt; + span?.setAttributes({ + 'workflow.stream.ws.write_to_send_ms': sentAt - write.startedAt, + ...(connection?.openedAt + ? { + 'workflow.stream.ws.open_to_send_ms': Math.max( + 0, + sentAt - connection.openedAt + ), + } + : {}), + }); +} + +function recordFirstWriteReply( + span: Span | undefined, + write: WriteTiming, + pending: PendingRequest | undefined +): void { + if (pending?.sentAt === undefined || pending.replyReceivedAt === undefined) { + return; + } + const processedAt = now(); + span?.setAttributes({ + 'workflow.stream.ws.send_to_reply_ms': + pending.replyReceivedAt - pending.sentAt, + 'workflow.stream.ws.reply_processing_ms': + processedAt - pending.replyReceivedAt, + 'workflow.stream.ws.write_total_ms': processedAt - write.startedAt, + }); +} const OIDC_FORCE_REFRESH_BUFFER_MS = 24 * 60 * 60 * 1000; function readAuthorization(headers: Headers): string | null { @@ -75,14 +186,7 @@ class VercelStreamWriteSession implements StreamWriteSession { private tail = Promise.resolve(); private inbound = Promise.resolve(); private nextReqId = 1; - private pending: - | { - reqId: number; - resolve: (meta: Record) => void; - reject: (error: unknown) => void; - timer: ReturnType; - } - | undefined; + private pending: PendingRequest | undefined; private poisonError: unknown; private wsUrl: string | undefined; private closeAcknowledged = false; @@ -91,6 +195,10 @@ class VercelStreamWriteSession implements StreamWriteSession { private drainTimer: ReturnType | undefined; private releaseDrainWait: (() => void) | undefined; private lastAuthorization: string | null = null; + private readonly sessionCreatedAt = now(); + private sessionHasWrite = false; + private connectionAttempt = 0; + private connectionTiming: ConnectionTiming | undefined; constructor( private readonly runId: string, @@ -107,29 +215,42 @@ class VercelStreamWriteSession implements StreamWriteSession { } write(chunkSeq: number, chunks: (string | Uint8Array)[]): Promise { - return this.enqueue(async () => { - this.assertUsable(); - await this.transportDecision; - this.assertUsable(); - if (this.mode === 'http') { - await this.writeHttp(chunks); - return; - } - // Core's default group cap equals the wire cap, so splitting is normally - // dormant. Keep it here as a guard against configured or future cap drift; - // v1 deliberately defines no separate whole-message byte budget. - for ( - let offset = 0; - offset < chunks.length; - offset += STREAM_WS_V1_MAX_CHUNKS_PER_WRITE - ) { - const batch = chunks.slice( - offset, - offset + STREAM_WS_V1_MAX_CHUNKS_PER_WRITE - ); - let reply: Record; - try { - reply = await this.request((reqId) => + const timing: WriteTiming = { + startedAt: now(), + sessionFirstWrite: !this.sessionHasWrite, + }; + this.sessionHasWrite = true; + return this.enqueue(() => this.writeInternal(chunkSeq, chunks, timing)); + } + + private async writeInternal( + chunkSeq: number, + chunks: (string | Uint8Array)[], + timing: WriteTiming + ): Promise { + this.assertUsable(); + await this.transportDecision; + this.assertUsable(); + if (this.mode === 'http') { + await this.writeHttp(chunks); + return; + } + // Core's default group cap equals the wire cap, so splitting is normally + // dormant. Keep it here as a guard against configured or future cap drift; + // v1 deliberately defines no separate whole-message byte budget. + for ( + let offset = 0; + offset < chunks.length; + offset += STREAM_WS_V1_MAX_CHUNKS_PER_WRITE + ) { + const batch = chunks.slice( + offset, + offset + STREAM_WS_V1_MAX_CHUNKS_PER_WRITE + ); + let reply: Record; + try { + reply = await this.request( + (reqId) => encodeStreamWsWriteRequest( { type: 'write', @@ -138,21 +259,21 @@ class VercelStreamWriteSession implements StreamWriteSession { numChunks: batch.length, }, batch - ) - ); - } catch (error) { - if (!(error instanceof StreamWsRequestNotSentError)) throw error; - this.fallbackToHttpBeforeSend(); - await this.writeHttp(chunks.slice(offset)); - return; - } - if (reply.type !== 'write_ack') { - throw this.poison( - new Error(`stream WebSocket write received ${reply.type}`) - ); - } + ), + offset === 0 ? timing : undefined + ); + } catch (error) { + if (!(error instanceof StreamWsRequestNotSentError)) throw error; + this.fallbackToHttpBeforeSend(); + await this.writeHttp(chunks.slice(offset)); + return; } - }); + if (reply.type !== 'write_ack') { + throw this.poison( + new Error(`stream WebSocket write received ${reply.type}`) + ); + } + } } dispose(): void { @@ -234,19 +355,31 @@ class VercelStreamWriteSession implements StreamWriteSession { } private startConnect(forceRefresh = false): Promise { - return this.connectSocket(forceRefresh).catch(() => { + const startedAt = now(); + const timing: ConnectionTiming = { + attempt: ++this.connectionAttempt, + configStartedAt: startedAt, + firstWriteSent: false, + }; + this.connectionTiming = timing; + return this.connectSocket(forceRefresh, timing).catch(() => { // Every failure before OPEN is a safe, session-long HTTP fallback. The // HTTP request itself still surfaces auth/configuration errors normally. if (this.mode === 'connecting') this.mode = 'http'; }); } - private async connectSocket(forceRefresh: boolean): Promise { + private async connectSocket( + forceRefresh: boolean, + timing: ConnectionTiming + ): Promise { if (!isWsStreamsTransportEnabled()) { this.mode = 'http'; return; } if (forceRefresh) { + // Included in config_token_ms: a slow refresh must not disappear from the + // first write after an auth-expiry reconnect. // Outside a Vercel function this invalidates @vercel/oidc's cached token. // Inside one, the invocation header remains authoritative; reconnecting // still re-resolves headers rather than retaining the old upgrade object. @@ -254,9 +387,13 @@ class VercelStreamWriteSession implements StreamWriteSession { expirationBufferMs: OIDC_FORCE_REFRESH_BUFFER_MS, }).catch(() => undefined); } + const httpPromise = getHttpConfig(this.config).then((http) => { + timing.configFinishedAt = now(); + return http; + }); const [{ WebSocket: WebSocketImpl }, http] = await Promise.all([ import('ws'), - getHttpConfig(this.config), + httpPromise, ]); if (this.mode !== 'connecting') return; if ( @@ -291,6 +428,7 @@ class VercelStreamWriteSession implements StreamWriteSession { async () => { await injectTraceContextIntoHeaders(http.headers); if (this.mode !== 'connecting') return; + timing.socketStartedAt = now(); const ws = new WebSocketImpl(url, { headers: headersToRecord(http.headers), }); @@ -306,6 +444,7 @@ class VercelStreamWriteSession implements StreamWriteSession { }; ws.once('open', () => { opened = true; + timing.openedAt = now(); if (this.mode !== 'connecting') { ws.close(1000, 'HTTP fallback selected'); resolve(); @@ -347,8 +486,9 @@ class VercelStreamWriteSession implements StreamWriteSession { ); }); ws.on('message', (raw) => { + const receivedAt = now(); this.inbound = this.inbound.then(() => - this.handleMessage(asBytes(raw)) + this.handleMessage(asBytes(raw), receivedAt) ); }); }); @@ -356,7 +496,10 @@ class VercelStreamWriteSession implements StreamWriteSession { ); } - private async handleMessage(raw: Uint8Array): Promise { + private async handleMessage( + raw: Uint8Array, + receivedAt: number + ): Promise { try { const frame = await decodeOne(raw); const reply = parseStreamWsReply(frame.meta, frame.body); @@ -378,6 +521,7 @@ class VercelStreamWriteSession implements StreamWriteSession { throw new Error('stream WebSocket reply cannot be correlated'); } this.pending = undefined; + pending.replyReceivedAt = receivedAt; clearTimeout(pending.timer); if (reply.type === 'close_ack') this.closeAcknowledged = true; if (reply.type === 'error') { @@ -500,7 +644,8 @@ class VercelStreamWriteSession implements StreamWriteSession { } private async request( - buildFrame: (reqId: number) => Uint8Array + buildFrame: (reqId: number) => Uint8Array, + writeTiming?: WriteTiming ): Promise> { this.assertUsable(); const ws = this.socket; @@ -516,6 +661,12 @@ class VercelStreamWriteSession implements StreamWriteSession { } catch (error) { throw new StreamWsRequestNotSentError(error); } + const connectionTiming = this.connectionTiming; + const connectionFirstWrite = connectionTiming?.firstWriteSent === false; + const detailedTiming = + writeTiming && (writeTiming.sessionFirstWrite || connectionFirstWrite) + ? writeTiming + : undefined; return withHttpClientSpan( { method: 'POST', @@ -524,31 +675,59 @@ class VercelStreamWriteSession implements StreamWriteSession { attributes: { 'workflow.stream.transport': 'ws', 'workflow.stream.ws.req_id': reqId, + ...firstWriteAttributes( + detailedTiming, + connectionTiming, + this.sessionCreatedAt + ), }, }, - async () => - new Promise>((resolve, reject) => { - // With no v1 progress/control reply, silence cannot distinguish a - // slow accepted request from a dead socket. Expiry is therefore an - // unknown outcome and must poison rather than replay. - const timer = setTimeout(() => { - this.failUnknown( - new Error( - `stream WebSocket request ${reqId} timed out with no reply` - ) - ); - }, getRequestTimeoutMs()); - timer.unref?.(); - this.pending = { reqId, resolve, reject, timer }; - try { - ws.send(frame, (error) => { - if (!error) return; + async (span) => { + if (detailedTiming) { + recordFirstWriteSetup(span, detailedTiming, connectionTiming); + } + let pending: PendingRequest | undefined; + const response = new Promise>( + (resolve, reject) => { + // With no v1 progress/control reply, silence cannot distinguish a + // slow accepted request from a dead socket. Expiry is therefore an + // unknown outcome and must poison rather than replay. + const timer = setTimeout(() => { + this.failUnknown( + new Error( + `stream WebSocket request ${reqId} timed out with no reply` + ) + ); + }, getRequestTimeoutMs()); + timer.unref?.(); + pending = { reqId, resolve, reject, timer }; + this.pending = pending; + try { + if (detailedTiming) { + recordFirstWriteSend( + span, + detailedTiming, + connectionTiming, + pending + ); + } + ws.send(frame, (error) => { + if (!error) return; + this.failUnknown(error); + }); + } catch (error) { this.failUnknown(error); - }); - } catch (error) { - this.failUnknown(error); + } } - }) + ); + try { + return await response; + } finally { + if (detailedTiming) { + recordFirstWriteReply(span, detailedTiming, pending); + } + } + } ); }