@xeonr/upload-pool-sdk 1.3.0 → 1.4.0

package/src/job-context.ts

@@ -7,6 +7,7 @@
  */
 import { writeFile } from "node:fs/promises";
 import { create } from "@bufbuild/protobuf";
+import { SpanStatusCode } from "@opentelemetry/api";
 import {
 	ConfirmMetaUploadRequestSchema,
 	InternalUpdateUploadRequestSchema,
@@ -18,6 +19,12 @@ import {
 import { UploadMetadataSchema } from "./protocol/uplim/api/v1/metadata_pb.js";
 import type { RpcClients } from "./rpc-clients.js";
 import type { Logger } from "./logger.js";
+import {
+	SPAN_ATTR,
+	getTracingHandle,
+	recordSpanError,
+	stampJobAttributes,
+} from "./tracing.js";
 import type {
 	FolderContext,
 	JobContext,
@@ -50,13 +57,82 @@ export interface JobEnvelope {
 		config?: string; // base64-encoded
 	};
 	folderContext?: FolderContext;
+	/**
+	 * W3C traceparent string forwarded by pipeline-worker. The SDK parses
+	 * this and uses it as the parent SpanContext for the iq.job root
+	 * span so worker callbacks nest under the dispatching parseUpload
+	 * trace. Optional — may be absent for legacy direct-dispatch paths.
+	 */
+	traceContext?: string;
+	/**
+	 * The pipeline workflow run id (workflow_runs.run_id) that owns this
+	 * job. The SDK stamps it on every span as the `pipeline.run_id`
+	 * attribute so pipeline-api's trace-store receiver associates the
+	 * span with the right run. Optional in payload; spans without it
+	 * are dropped by the receiver.
+	 */
+	wfRunId?: string;
+}
+
+/**
+ * Wrap an async fn in an `iq.ctx.*` span stamped with the job's identifiers.
+ * Used for every ctx.* method so the admin UI sees a row per handler
+ * callback with its own latency and (where applicable) request/response
+ * body on span events.
+ *
+ * No-op when tracing is uninitialised (e.g. in tests): the inner fn runs
+ * directly and we lose the span, no exception.
+ */
+function withCtxSpan<T>(
+	name: string,
+	envelope: JobEnvelope,
+	workerId: string | undefined,
+	fn: () => Promise<T>,
+	attrs?: Record<string, string | number | boolean>,
+): Promise<T> {
+	const handle = getTracingHandle();
+	if (!handle) {
+		return fn();
+	}
+	return handle.tracer.startActiveSpan(name, async (span) => {
+		stampJobAttributes(span, {
+			runId: envelope.wfRunId ?? "",
+			jobId: envelope.jobId,
+			uploadId: envelope.uploadId,
+			urn: envelope.contentTypeContext.urn,
+			workerId,
+		});
+		if (attrs) {
+			for (const [k, v] of Object.entries(attrs)) {
+				span.setAttribute(k, v);
+			}
+		}
+		const startedAt = Date.now();
+		try {
+			const result = await fn();
+			span.setAttribute("ctx.duration_ms", Date.now() - startedAt);
+			span.setStatus({ code: SpanStatusCode.OK });
+			return result;
+		} catch (err) {
+			recordSpanError(span, err);
+			throw err;
+		} finally {
+			span.end();
+		}
+	});
+}
+
+export interface CreateJobContextOpts {
+	workerId?: string;
 }
 
 export function createJobContext(
 	envelope: JobEnvelope,
 	rpc: RpcClients,
 	logger: Logger,
+	opts: CreateJobContextOpts = {},
 ): JobContext {
+	const workerId = opts.workerId;
 	const ctx: JobContext = {
 		jobId: envelope.jobId,
 		uploadId: envelope.uploadId,
@@ -73,196 +149,300 @@ export function createJobContext(
 		sourceUrl: envelope.sourceUrl,
 
 		async download() {
-			logger.debug("source.download.start");
-			const resp = await fetch(envelope.sourceUrl);
-			if (!resp.ok || !resp.body) {
-				logger.error("source.download.failed", {
-					status: resp.status,
-					statusText: resp.statusText,
-				});
-				throw new Error(`source download failed: ${resp.status} ${resp.statusText}`);
-			}
-			return resp.body;
+			return withCtxSpan("iq.ctx.download", envelope, workerId, async () => {
+				logger.debug("source.download.start");
+				const resp = await fetch(envelope.sourceUrl);
+				if (!resp.ok || !resp.body) {
+					logger.error("source.download.failed", {
+						status: resp.status,
+						statusText: resp.statusText,
+					});
+					throw new Error(
+						`source download failed: ${resp.status} ${resp.statusText}`,
+					);
+				}
+				return resp.body;
+			}, { "ctx.mode": "stream" });
 		},
 
 		async downloadBuffer() {
-			logger.debug("source.download.start", { mode: "buffer" });
-			const resp = await fetch(envelope.sourceUrl);
-			if (!resp.ok) {
-				logger.error("source.download.failed", {
-					status: resp.status,
-					statusText: resp.statusText,
-				});
-				throw new Error(`source download failed: ${resp.status} ${resp.statusText}`);
-			}
-			const buf = await resp.arrayBuffer();
-			logger.debug("source.download.complete", { sizeBytes: buf.byteLength });
-			return new Uint8Array(buf);
+			return withCtxSpan("iq.ctx.download", envelope, workerId, async () => {
+				logger.debug("source.download.start", { mode: "buffer" });
+				const resp = await fetch(envelope.sourceUrl);
+				if (!resp.ok) {
+					logger.error("source.download.failed", {
+						status: resp.status,
+						statusText: resp.statusText,
+					});
+					throw new Error(
+						`source download failed: ${resp.status} ${resp.statusText}`,
+					);
+				}
+				const buf = await resp.arrayBuffer();
+				logger.debug("source.download.complete", { sizeBytes: buf.byteLength });
+				return new Uint8Array(buf);
+			}, { "ctx.mode": "buffer" });
 		},
 
 		async downloadToFile(path: string) {
-			const data = await ctx.downloadBuffer();
-			await writeFile(path, data);
-			logger.debug("source.download.to_file", { path, sizeBytes: data.byteLength });
+			return withCtxSpan(
+				"iq.ctx.download_to_file",
+				envelope,
+				workerId,
+				async () => {
+					const data = await ctx.downloadBuffer();
+					await writeFile(path, data);
+					logger.debug("source.download.to_file", {
+						path,
+						sizeBytes: data.byteLength,
+					});
+				},
+			);
 		},
 
 		async uploadMeta(opts: UploadMetaOpts) {
-			logger.debug("meta.upload.start", {
-				type: opts.type,
-				instance: opts.instance,
-				filename: opts.filename,
-			});
-			const protoType = META_TYPE_MAP[opts.type];
-			const metaType = create(MetaTypeSchema, {
-				type: protoType,
-				n: opts.instance,
-				filename: opts.filename,
-			});
+			return withCtxSpan(
+				"iq.ctx.upload_meta",
+				envelope,
+				workerId,
+				async () => {
+					logger.debug("meta.upload.start", {
+						type: opts.type,
+						instance: opts.instance,
+						filename: opts.filename,
+					});
+					const protoType = META_TYPE_MAP[opts.type];
+					const metaType = create(MetaTypeSchema, {
+						type: protoType,
+						n: opts.instance,
+						filename: opts.filename,
+					});
 
-			// 1. RequestMetaUpload → presigned URL.
-			const requested = await rpc.internalUploads.requestMetaUpload(
-				create(RequestMetaUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					metaType: [metaType],
-				}),
-			);
+					// 1. RequestMetaUpload → presigned URL.
+					const requested = await rpc.internalUploads.requestMetaUpload(
+						create(RequestMetaUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							metaType: [metaType],
+						}),
+					);
 
-			const meta = requested.metaUploads[0];
-			if (!meta) {
-				throw new Error("RequestMetaUpload returned no meta_uploads");
-			}
+					const meta = requested.metaUploads[0];
+					if (!meta) {
+						throw new Error("RequestMetaUpload returned no meta_uploads");
+					}
 
-			// 2. S3 PUT.
-			const contentType = opts.contentType ?? inferContentType(opts.type);
-			const body = opts.data instanceof Uint8Array
-				? opts.data
-				: opts.data;
-			const putResp = await fetch(meta.uploadUrl, {
-				method: "PUT",
-				body: body as BodyInit,
-				headers: { "Content-Type": contentType },
-				duplex: "half",
-			} as RequestInit & { duplex?: "half" });
-			if (!putResp.ok) {
-				throw new Error(`meta upload PUT failed: ${putResp.status} ${putResp.statusText}`);
-			}
+					// 2. S3 PUT — own span so the upload's bytes/duration are
+					// visible separately from the API round-trips.
+					const contentType =
+						opts.contentType ?? inferContentType(opts.type);
+					const body =
+						opts.data instanceof Uint8Array ? opts.data : opts.data;
+					await withCtxSpan(
+						"iq.s3.put",
+						envelope,
+						workerId,
+						async () => {
+							const handle = getTracingHandle();
+							const span = handle?.tracer
+								? undefined
+								: undefined;
+							void span;
+							const putResp = await fetch(meta.uploadUrl, {
+								method: "PUT",
+								body: body as BodyInit,
+								headers: { "Content-Type": contentType },
+								duplex: "half",
+							} as RequestInit & { duplex?: "half" });
+							if (!putResp.ok) {
+								throw new Error(
+									`meta upload PUT failed: ${putResp.status} ${putResp.statusText}`,
+								);
+							}
+						},
+						{
+							[SPAN_ATTR.S3_URL_ORIGIN]: safeOrigin(meta.uploadUrl),
+							[SPAN_ATTR.S3_BYTES]:
+								body instanceof Uint8Array ? body.byteLength : 0,
+							"s3.content_type": contentType,
+							"s3.meta_type": opts.type,
+						},
+					);
 
-			// 3. ConfirmMetaUpload.
-			await rpc.internalUploads.confirmMetaUpload(
-				create(ConfirmMetaUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					metadataId: [meta.metadataId],
-				}),
+					// 3. ConfirmMetaUpload.
+					await rpc.internalUploads.confirmMetaUpload(
+						create(ConfirmMetaUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							metadataId: [meta.metadataId],
+						}),
+					);
+					logger.info("meta.upload.complete", {
+						type: opts.type,
+						metadataId: meta.metadataId,
+					});
+				},
+				{ "ctx.meta_type": opts.type },
 			);
-			logger.info("meta.upload.complete", {
-				type: opts.type,
-				metadataId: meta.metadataId,
-			});
 		},
 
 		async setMetadata(metadata: UploadMetadataPartial | UploadMetadataProto) {
-			// Distinguish a pre-built proto from a partial: protos carry the
-			// @bufbuild/protobuf $typeName brand.
-			const isProto = (metadata as UploadMetadataProto).$typeName !== undefined;
-			let protoMetadata: ReturnType<typeof create<typeof UploadMetadataSchema>>;
-			let thumbnailVersion: number | undefined;
-			if (isProto) {
-				protoMetadata = metadata as ReturnType<typeof create<typeof UploadMetadataSchema>>;
-			} else {
-				const partial = metadata as UploadMetadataPartial;
-				protoMetadata = create(UploadMetadataSchema, {});
-				// Server-side proto has typed sub-messages; the SDK's
-				// UploadMetadataPartial is intentionally loose so workers can
-				// populate just what they computed. Cast into the proto via
-				// the JSON path — fields that don't map are ignored.
-				Object.assign(protoMetadata as object, partial);
-				thumbnailVersion = partial.thumbnailGenerationVersion;
-			}
+			return withCtxSpan(
+				"iq.ctx.set_metadata",
+				envelope,
+				workerId,
+				async () => {
+					// Distinguish a pre-built proto from a partial: protos carry
+					// the @bufbuild/protobuf $typeName brand.
+					const isProto =
+						(metadata as UploadMetadataProto).$typeName !== undefined;
+					let protoMetadata: ReturnType<
+						typeof create<typeof UploadMetadataSchema>
+					>;
+					let thumbnailVersion: number | undefined;
+					if (isProto) {
+						protoMetadata = metadata as ReturnType<
+							typeof create<typeof UploadMetadataSchema>
+						>;
+					} else {
+						const partial = metadata as UploadMetadataPartial;
+						protoMetadata = create(UploadMetadataSchema, {});
+						// Server-side proto has typed sub-messages; the SDK's
+						// UploadMetadataPartial is intentionally loose so workers
+						// can populate just what they computed. Cast into the
+						// proto via the JSON path — fields that don't map are
+						// ignored.
+						Object.assign(protoMetadata as object, partial);
+						thumbnailVersion = partial.thumbnailGenerationVersion;
+					}
 
-			await rpc.internalUploads.updateUpload(
-				create(InternalUpdateUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					metadata: protoMetadata,
-					thumbnailGenerationVersion: thumbnailVersion,
-				}),
+					await rpc.internalUploads.updateUpload(
+						create(InternalUpdateUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							metadata: protoMetadata,
+							thumbnailGenerationVersion: thumbnailVersion,
+						}),
+					);
+					logger.debug("metadata.set", { thumbnailVersion });
+				},
 			);
-			logger.debug("metadata.set", { thumbnailVersion });
 		},
 
 		async markHasThumbnail(version: number) {
-			await rpc.internalUploads.updateUpload(
-				create(InternalUpdateUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					hasThumbnail: true,
-					thumbnailGenerationVersion: version,
-				}),
+			return withCtxSpan(
+				"iq.ctx.mark_has_thumbnail",
+				envelope,
+				workerId,
+				async () => {
+					await rpc.internalUploads.updateUpload(
+						create(InternalUpdateUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							hasThumbnail: true,
+							thumbnailGenerationVersion: version,
+						}),
+					);
+					logger.info("thumbnail.marked", { version });
+				},
+				{ "ctx.thumbnail_version": version },
 			);
-			logger.info("thumbnail.marked", { version });
 		},
 
 		async setDescription(text: string, tags?: string[]) {
-			await rpc.internalUploads.updateUpload(
-				create(InternalUpdateUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					aiMetadata: create(InternalUpdateUploadRequest_AiMetadataSchema, {
-						description: text,
-						generatedTags: tags ?? [],
-					}),
-				}),
+			return withCtxSpan(
+				"iq.ctx.set_description",
+				envelope,
+				workerId,
+				async () => {
+					await rpc.internalUploads.updateUpload(
+						create(InternalUpdateUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							aiMetadata: create(
+								InternalUpdateUploadRequest_AiMetadataSchema,
+								{
+									description: text,
+									generatedTags: tags ?? [],
+								},
+							),
+						}),
+					);
+					logger.info("description.set", {
+						length: text.length,
+						tags: tags?.length ?? 0,
+					});
+				},
+				{
+					"ctx.description_length": text.length,
+					"ctx.tag_count": tags?.length ?? 0,
+				},
 			);
-			logger.info("description.set", {
-				length: text.length,
-				tags: tags?.length ?? 0,
-			});
 		},
 
 		async withPresignedUploads<T>(
 			opts: PresignedUploadRequest[],
 			callback: (handles: import("./types.js").MetaUploadHandle[]) => Promise<T>,
 		): Promise<T> {
-			logger.debug("presigned.upload.request", {
-				count: opts.length,
-				types: opts.map((o) => o.type),
-			});
-			const metaTypes = opts.map((opt) => create(MetaTypeSchema, {
-				type: META_TYPE_MAP[opt.type],
-				n: opt.instance,
-				filename: opt.filename,
-			}));
-			const requested = await rpc.internalUploads.requestMetaUpload(
-				create(RequestMetaUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					metaType: metaTypes,
-				}),
+			return withCtxSpan(
+				"iq.ctx.with_presigned_uploads",
+				envelope,
+				workerId,
+				async () => {
+					logger.debug("presigned.upload.request", {
+						count: opts.length,
+						types: opts.map((o) => o.type),
+					});
+					const metaTypes = opts.map((opt) =>
+						create(MetaTypeSchema, {
+							type: META_TYPE_MAP[opt.type],
+							n: opt.instance,
+							filename: opt.filename,
+						}),
+					);
+					const requested = await rpc.internalUploads.requestMetaUpload(
+						create(RequestMetaUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							metaType: metaTypes,
+						}),
+					);
+					if (requested.metaUploads.length !== opts.length) {
+						throw new Error(
+							`RequestMetaUpload returned ${requested.metaUploads.length} URLs, expected ${opts.length}`,
+						);
+					}
+					const handles = requested.metaUploads.map((m) => ({
+						uploadUrl: m.uploadUrl,
+						downloadUrl: m.downloadUrl,
+						metadataId: m.metadataId,
+					}));
+					const result = await callback(handles);
+					// Confirm only after the callback resolves cleanly — on
+					// throw, the rows stay unfinalised and will be cleaned up
+					// by the next confirm sweep or expire naturally.
+					await rpc.internalUploads.confirmMetaUpload(
+						create(ConfirmMetaUploadRequestSchema, {
+							updateToken: envelope.updateToken,
+							metadataId: requested.metaUploads.map((m) => m.metadataId),
+						}),
+					);
+					logger.info("presigned.upload.complete", {
+						count: handles.length,
+					});
+					return result;
+				},
+				{ "ctx.upload_count": opts.length },
 			);
-			if (requested.metaUploads.length !== opts.length) {
-				throw new Error(
-					`RequestMetaUpload returned ${requested.metaUploads.length} URLs, expected ${opts.length}`,
-				);
-			}
-			const handles = requested.metaUploads.map((m) => ({
-				uploadUrl: m.uploadUrl,
-				downloadUrl: m.downloadUrl,
-				metadataId: m.metadataId,
-			}));
-			const result = await callback(handles);
-			// Confirm only after the callback resolves cleanly — on throw,
-			// the rows stay unfinalised and will be cleaned up by the next
-			// confirm sweep or expire naturally.
-			await rpc.internalUploads.confirmMetaUpload(
-				create(ConfirmMetaUploadRequestSchema, {
-					updateToken: envelope.updateToken,
-					metadataId: requested.metaUploads.map((m) => m.metadataId),
-				}),
-			);
-			logger.info("presigned.upload.complete", { count: handles.length });
-			return result;
 		},
 	};
 
 	return ctx;
 }
 
+function safeOrigin(url: string): string {
+	try {
+		const u = new URL(url);
+		return `${u.protocol}//${u.host}${u.pathname}`;
+	} catch {
+		return "<malformed>";
+	}
+}
+
 function inferContentType(metaType: UploadMetaOpts["type"]): string {
 	switch (metaType) {
 		case "THUMBNAIL_LIGHT":