diff --git a/src/common/otel-trace.ts b/src/common/otel-trace.ts index 9a74734..3204e91 100644 --- a/src/common/otel-trace.ts +++ b/src/common/otel-trace.ts @@ -1,9 +1,42 @@ import { createHash, randomBytes } from "node:crypto"; +import { AgentRunError } from "./errors.js"; import type { CommandRecord, JsonRecord, RunRecord } from "./types.js"; const OTLP_TIMEOUT_MS = 1500; const ZERO_TRACE_ID = "00000000000000000000000000000000"; const ZERO_SPAN_ID = "0000000000000000"; +const MAX_DIAGNOSTIC_OPERATION_LENGTH = 128; +const MAX_DIAGNOSTIC_CORRELATION_ID_LENGTH = 256; + +export const AGENTRUN_TRACE_PURPOSE_HEADER = "x-agentrun-trace-purpose"; +export const AGENTRUN_DIAGNOSTIC_OPERATION_HEADER = "x-agentrun-diagnostic-operation"; +export const AGENTRUN_DIAGNOSTIC_TARGET_BUSINESS_TRACE_ID_HEADER = "x-agentrun-diagnostic-target-business-trace-id"; +export const AGENTRUN_DIAGNOSTIC_TARGET_COMMAND_ID_HEADER = "x-agentrun-diagnostic-target-command-id"; + +export interface AgentRunBusinessOtelTraceContext { + kind: "business"; + businessTraceId: string; + traceId: string; + parentSpanId: string; + traceparent: string; + valuesPrinted: false; +} + +export interface AgentRunDiagnosticOtelTraceContext { + kind: "diagnostic"; + traceId: string; + parentSpanId: string; + traceparent: string; + operation: string; + targetBusinessTraceId: string; + targetCommandId: string; + targetOtelTraceId: string; + targetSpanId: string; + valuesPrinted: false; +} + +export type AgentRunOtelTraceContext = AgentRunBusinessOtelTraceContext | AgentRunDiagnosticOtelTraceContext; +export type AgentRunOtelRequestHeaders = Record; export function agentRunBusinessTraceId(run: RunRecord | null | undefined, command?: CommandRecord | null): string | null { const traceSink = asRecord(run?.traceSink); @@ -22,17 +55,62 @@ export function agentRunBusinessTraceId(run: RunRecord | null | undefined, comma return null; } -export function agentRunOtelTraceContext(run: RunRecord | null | undefined, command?: CommandRecord | null) { +export function agentRunOtelTraceContext(run: RunRecord | null | undefined, command?: CommandRecord | null): AgentRunBusinessOtelTraceContext { const businessTraceId = agentRunBusinessTraceId(run, command) ?? `run_${String(run?.id ?? "unknown")}`; - const traceId = nonZeroHex(createHash("sha256").update(`hwlab-code-agent:${businessTraceId}`).digest("hex").slice(0, 32), ZERO_TRACE_ID); - const parentSpanId = nonZeroHex(createHash("sha256").update(`agentrun-manager-parent:${businessTraceId}`).digest("hex").slice(0, 16), ZERO_SPAN_ID); - return { businessTraceId, traceId, parentSpanId, traceparent: `00-${traceId}-${parentSpanId}-01`, valuesPrinted: false }; + return agentRunBusinessOtelTraceContext(businessTraceId); } -export async function emitAgentRunOtelSpan(name: string, run: RunRecord | null | undefined, env: NodeJS.ProcessEnv = process.env, options: { command?: CommandRecord | null; startTimeMs?: number; endTimeMs?: number; status?: "ok" | "error"; error?: unknown; attributes?: JsonRecord; kind?: number; scopeName?: string } = {}): Promise { +export function agentRunBusinessOtelTraceContext(businessTraceId: string): AgentRunBusinessOtelTraceContext { + const traceId = nonZeroHex(createHash("sha256").update(`hwlab-code-agent:${businessTraceId}`).digest("hex").slice(0, 32), ZERO_TRACE_ID); + const parentSpanId = nonZeroHex(createHash("sha256").update(`agentrun-manager-parent:${businessTraceId}`).digest("hex").slice(0, 16), ZERO_SPAN_ID); + return { kind: "business", businessTraceId, traceId, parentSpanId, traceparent: `00-${traceId}-${parentSpanId}-01`, valuesPrinted: false }; +} + +export function agentRunDiagnosticOtelTraceContext(headers: AgentRunOtelRequestHeaders): AgentRunDiagnosticOtelTraceContext | null { + const purpose = singleHeader(headers, AGENTRUN_TRACE_PURPOSE_HEADER); + const operation = singleHeader(headers, AGENTRUN_DIAGNOSTIC_OPERATION_HEADER); + const targetBusinessTraceId = singleHeader(headers, AGENTRUN_DIAGNOSTIC_TARGET_BUSINESS_TRACE_ID_HEADER); + const targetCommandId = singleHeader(headers, AGENTRUN_DIAGNOSTIC_TARGET_COMMAND_ID_HEADER); + if (purpose === null) { + if (operation !== null || targetBusinessTraceId !== null || targetCommandId !== null) { + throw diagnosticContextError(`${AGENTRUN_TRACE_PURPOSE_HEADER}: diagnostic is required when diagnostic correlation headers are present`); + } + return null; + } + if (purpose !== "diagnostic") throw diagnosticContextError(`${AGENTRUN_TRACE_PURPOSE_HEADER} must be diagnostic`); + const traceparent = requiredSingleHeader(headers, "traceparent"); + const parsedTraceparent = parseTraceparent(traceparent); + if (operation === null || operation.length > MAX_DIAGNOSTIC_OPERATION_LENGTH || !/^[A-Za-z0-9][A-Za-z0-9._:/-]*$/u.test(operation)) { + throw diagnosticContextError(`${AGENTRUN_DIAGNOSTIC_OPERATION_HEADER} must be a bounded semantic operation`); + } + if (targetBusinessTraceId === null || targetBusinessTraceId.length > MAX_DIAGNOSTIC_CORRELATION_ID_LENGTH || !/^trc_[A-Za-z0-9_.:-]+$/u.test(targetBusinessTraceId)) { + throw diagnosticContextError(`${AGENTRUN_DIAGNOSTIC_TARGET_BUSINESS_TRACE_ID_HEADER} must be a business trace id`); + } + if (targetCommandId === null || targetCommandId.length > MAX_DIAGNOSTIC_CORRELATION_ID_LENGTH || !/^cmd_[A-Za-z0-9_.:-]+$/u.test(targetCommandId)) { + throw diagnosticContextError(`${AGENTRUN_DIAGNOSTIC_TARGET_COMMAND_ID_HEADER} must be a command id`); + } + const target = agentRunBusinessOtelTraceContext(targetBusinessTraceId); + if (parsedTraceparent.traceId === target.traceId) { + throw diagnosticContextError("diagnostic traceId must differ from the target business OTel traceId"); + } + return { + kind: "diagnostic", + traceId: parsedTraceparent.traceId, + parentSpanId: parsedTraceparent.parentSpanId, + traceparent, + operation, + targetBusinessTraceId, + targetCommandId, + targetOtelTraceId: target.traceId, + targetSpanId: target.parentSpanId, + valuesPrinted: false, + }; +} + +export async function emitAgentRunOtelSpan(name: string, run: RunRecord | null | undefined, env: NodeJS.ProcessEnv = process.env, options: { command?: CommandRecord | null; traceContext?: AgentRunOtelTraceContext; startTimeMs?: number; endTimeMs?: number; status?: "ok" | "error"; error?: unknown; attributes?: JsonRecord; kind?: number; scopeName?: string } = {}): Promise { const endpoint = resolveOtlpTracesEndpoint(env); if (!endpoint || typeof fetch !== "function") return { ok: false, skipped: true, reason: "otlp-endpoint-missing", valuesPrinted: false }; - const context = agentRunOtelTraceContext(run, options.command ?? null); + const context = options.traceContext ?? agentRunOtelTraceContext(run, options.command ?? null); const startedAtMs = Number.isFinite(Number(options.startTimeMs)) ? Number(options.startTimeMs) : Date.now(); const endedAtMs = Number.isFinite(Number(options.endTimeMs)) ? Number(options.endTimeMs) : Date.now(); const spanId = nonZeroHex(randomBytes(8).toString("hex"), ZERO_SPAN_ID); @@ -52,15 +130,15 @@ export async function emitAgentRunOtelSpan(name: string, run: RunRecord | null | startTimeUnixNano: unixNano(startedAtMs), endTimeUnixNano: unixNano(Math.max(startedAtMs, endedAtMs)), attributes: attributesFromRecord({ - traceId: context.businessTraceId, - "otel.trace_id": context.traceId, "agentrun.stage": name, runId: run?.id ?? null, commandId: options.command?.id ?? null, sessionId: run?.sessionRef?.sessionId ?? null, ...options.attributes, + ...spanCorrelationAttributes(context), valuesPrinted: false, }), + ...(context.kind === "diagnostic" ? { links: [diagnosticSpanLink(context)] } : {}), status: { code: statusCode, ...(options.error ? { message: String((options.error as Error)?.message ?? options.error).slice(0, 300) } : {}), @@ -86,6 +164,72 @@ export async function emitAgentRunOtelSpan(name: string, run: RunRecord | null | } } +function spanCorrelationAttributes(context: AgentRunOtelTraceContext): JsonRecord { + if (context.kind === "business") return { + traceId: context.businessTraceId, + "otel.trace_id": context.traceId, + }; + return { + traceId: null, + "otel.trace_id": context.traceId, + "diagnostic.target_business_trace_id": context.targetBusinessTraceId, + "diagnostic.target_otel_trace_id": context.targetOtelTraceId, + "diagnostic.operation": context.operation, + commandId: context.targetCommandId, + }; +} + +function diagnosticSpanLink(context: AgentRunDiagnosticOtelTraceContext): JsonRecord { + return { + traceId: context.targetOtelTraceId, + spanId: context.targetSpanId, + attributes: attributesFromRecord({ + "diagnostic.target_business_trace_id": context.targetBusinessTraceId, + "diagnostic.target_otel_trace_id": context.targetOtelTraceId, + "diagnostic.operation": context.operation, + commandId: context.targetCommandId, + }), + droppedAttributesCount: 0, + }; +} + +function parseTraceparent(value: string): { traceId: string; parentSpanId: string } { + const match = /^00-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/u.exec(value); + if (!match || match[1] === ZERO_TRACE_ID || match[2] === ZERO_SPAN_ID) { + throw diagnosticContextError("traceparent must be a valid sampled or unsampled W3C version 00 context"); + } + return { traceId: match[1], parentSpanId: match[2] }; +} + +function requiredSingleHeader(headers: AgentRunOtelRequestHeaders, name: string): string { + const value = singleHeader(headers, name); + if (value === null) throw diagnosticContextError(`${name} is required for diagnostic tracing`); + return value; +} + +function singleHeader(headers: AgentRunOtelRequestHeaders, name: string): string | null { + const entry = Object.entries(headers).find(([key]) => key.toLowerCase() === name); + if (entry === undefined) return null; + const value = entry[1]; + if (value === undefined) throw diagnosticContextError(`${name} must not be empty`); + if (Array.isArray(value)) { + if (value.length !== 1) throw diagnosticContextError(`${name} must have exactly one value`); + const normalized = value[0]?.trim() ?? ""; + if (normalized.length === 0) throw diagnosticContextError(`${name} must not be empty`); + return normalized; + } + const normalized = value.trim(); + if (normalized.length === 0) throw diagnosticContextError(`${name} must not be empty`); + return normalized; +} + +function diagnosticContextError(message: string): AgentRunError { + return new AgentRunError("schema-invalid", message, { + httpStatus: 400, + details: { reason: "diagnostic-trace-context-invalid", valuesPrinted: false }, + }); +} + function resolveOtlpTracesEndpoint(env: NodeJS.ProcessEnv): string | null { const explicit = firstNonEmpty(env.AGENTRUN_OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT); if (explicit) return explicit.replace(/\/+$/u, ""); diff --git a/src/mgr/server.ts b/src/mgr/server.ts index 74656bd..a478d46 100644 --- a/src/mgr/server.ts +++ b/src/mgr/server.ts @@ -24,7 +24,7 @@ import { getProviderProfileConfig, getProviderProfileValidation, listBackendCapa import { listToolCredentials, setGithubSshToolCredential, showToolCredential } from "./tool-credentials.js"; import { aipodSpecFromInput, applyAipodSpec, deleteAipodSpec, listAipodSpecs, renderAipodSpecByName, showAipodSpec } from "../common/aipod-specs.js"; import { staticWorkReadyCapabilitySummary } from "../common/work-ready.js"; -import { emitAgentRunOtelSpan } from "../common/otel-trace.js"; +import { agentRunBusinessTraceId, agentRunDiagnosticOtelTraceContext, emitAgentRunOtelSpan, type AgentRunDiagnosticOtelTraceContext } from "../common/otel-trace.js"; function pvcOptions(defaults: { kubectlCommand?: string } | undefined): SessionPvcOptions { return defaults?.kubectlCommand ? { kubectlCommand: defaults.kubectlCommand } : {}; @@ -271,14 +271,37 @@ export async function startManagerServer(options: ManagerServerOptions = {}): Pr const stopKafkaOutboxRelay = kafkaOutboxRelayOptions.enabled ? startKafkaEventOutboxRelay({ store, options: kafkaOutboxRelayOptions, onError: (error) => console.warn(JSON.stringify({ code: "kafka-outbox-relay-failed", message: error instanceof Error ? error.message : String(error), valuesPrinted: false })) }) : null; const server = createServer(async (req, res) => { const traceId = `trc_${Date.now().toString(36)}_${Math.random().toString(36).slice(2, 8)}`; + const method = req.method ?? "GET"; + const url = new URL(req.url ?? "/", "http://agentrun.local"); + let diagnosticTraceContext: AgentRunDiagnosticOtelTraceContext | null = null; try { - const method = req.method ?? "GET"; - const url = new URL(req.url ?? "/", "http://agentrun.local"); assertManagerRequestAuthorized(req, url.pathname, auth); - const data = await route({ method, url, body: await readBody(req), store, sourceCommit, authSummary, runnerReconcilerOptions, runnerDispatcherOptions, kafkaOutboxRelayOptions, resolvedRunnerJobDefaults, ...(runnerJobDefaults ? { runnerJobDefaults } : {}), ...(sessionPvcDefaults ? { sessionPvcDefaults } : {}), ...(providerProfileDefaults ? { providerProfileDefaults } : {}), ...(toolCredentialDefaults ? { toolCredentialDefaults } : {}), ...(aipodSpecDir ? { aipodSpecDir } : {}) }); + diagnosticTraceContext = agentRunDiagnosticOtelTraceContext(req.headers); + if (diagnosticTraceContext !== null && method !== "GET") { + throw new AgentRunError("schema-invalid", "diagnostic trace context is only valid for GET requests", { httpStatus: 400, details: { reason: "diagnostic-trace-context-write-denied", valuesPrinted: false } }); + } + if (diagnosticTraceContext !== null && !isSupportedDiagnosticReadPath(url.pathname)) { + throw new AgentRunError("schema-invalid", "diagnostic trace context is not supported for this GET route", { httpStatus: 400, details: { reason: "diagnostic-trace-route-unsupported", route: url.pathname, valuesPrinted: false } }); + } + const data = await route({ method, url, body: await readBody(req), store, sourceCommit, authSummary, diagnosticTraceContext, runnerReconcilerOptions, runnerDispatcherOptions, kafkaOutboxRelayOptions, resolvedRunnerJobDefaults, ...(runnerJobDefaults ? { runnerJobDefaults } : {}), ...(sessionPvcDefaults ? { sessionPvcDefaults } : {}), ...(providerProfileDefaults ? { providerProfileDefaults } : {}), ...(toolCredentialDefaults ? { toolCredentialDefaults } : {}), ...(aipodSpecDir ? { aipodSpecDir } : {}) }); writeJson(res, 200, { ok: true, data, traceId }); } catch (error) { const agentError = normalizeError(error); + if (diagnosticTraceContext !== null) { + void emitAgentRunOtelSpan("diagnostic_request_rejected", null, process.env, { + traceContext: diagnosticTraceContext, + kind: 2, + status: "error", + error: agentError, + attributes: { + "http.method": method, + "http.route": url.pathname, + "http.status_code": agentError.httpStatus, + failureKind: agentError.failureKind, + "diagnostic.rejection_reason": diagnosticRejectionReason(agentError), + }, + }); + } writeJson(res, agentError.httpStatus, { ok: false, failureKind: agentError.failureKind, message: agentError.message, traceId, error: errorToJson(error) }); } }); @@ -301,6 +324,41 @@ async function readBody(req: import("node:http").IncomingMessage): Promise { + if (input.context === null) return null; + if (input.command && input.command.id !== input.context.targetCommandId) { + throw diagnosticTargetMismatch("diagnostic target command does not match the requested command"); + } + const command = input.command ?? await input.store.getCommand(input.context.targetCommandId); + if (command.runId !== input.run.id) throw diagnosticTargetMismatch("diagnostic target command does not belong to the requested run"); + const businessTraceId = agentRunBusinessTraceId(input.run, command); + if (businessTraceId !== input.context.targetBusinessTraceId) { + throw diagnosticTargetMismatch("diagnostic target business trace does not match the requested run and command"); + } + return input.context; +} + +function diagnosticTargetMismatch(message: string): AgentRunError { + return new AgentRunError("schema-invalid", message, { + httpStatus: 400, + details: { reason: "diagnostic-trace-target-mismatch", valuesPrinted: false }, + }); +} + +function diagnosticRejectionReason(error: AgentRunError): string { + return typeof error.details?.reason === "string" ? error.details.reason : "diagnostic-request-rejected"; +} + async function refreshQueueTaskForRead(store: AgentRunStore, taskId: string): Promise { const task = await store.getQueueTask(taskId); const refreshed = await refreshQueueTaskRecordForRead(store, task); @@ -582,7 +640,7 @@ function commandResultOtelAttributes(result: JsonValue): JsonRecord { }; } -async function route({ method, url, body, store, sourceCommit, authSummary, runnerJobDefaults, sessionPvcDefaults, providerProfileDefaults, toolCredentialDefaults, runnerReconcilerOptions, runnerDispatcherOptions, kafkaOutboxRelayOptions, resolvedRunnerJobDefaults, aipodSpecDir }: { method: string; url: URL; body: unknown; store: AgentRunStore; sourceCommit: string; authSummary?: JsonRecord; runnerJobDefaults?: NonNullable; sessionPvcDefaults?: NonNullable; providerProfileDefaults?: NonNullable; toolCredentialDefaults?: NonNullable; runnerReconcilerOptions: RunnerReconcilerRuntimeOptions; runnerDispatcherOptions: RunnerDispatcherOptions; kafkaOutboxRelayOptions: KafkaOutboxRelayOptions; resolvedRunnerJobDefaults: RunnerJobDefaults; aipodSpecDir?: string }): Promise { +async function route({ method, url, body, store, sourceCommit, authSummary, diagnosticTraceContext, runnerJobDefaults, sessionPvcDefaults, providerProfileDefaults, toolCredentialDefaults, runnerReconcilerOptions, runnerDispatcherOptions, kafkaOutboxRelayOptions, resolvedRunnerJobDefaults, aipodSpecDir }: { method: string; url: URL; body: unknown; store: AgentRunStore; sourceCommit: string; authSummary?: JsonRecord; diagnosticTraceContext: AgentRunDiagnosticOtelTraceContext | null; runnerJobDefaults?: NonNullable; sessionPvcDefaults?: NonNullable; providerProfileDefaults?: NonNullable; toolCredentialDefaults?: NonNullable; runnerReconcilerOptions: RunnerReconcilerRuntimeOptions; runnerDispatcherOptions: RunnerDispatcherOptions; kafkaOutboxRelayOptions: KafkaOutboxRelayOptions; resolvedRunnerJobDefaults: RunnerJobDefaults; aipodSpecDir?: string }): Promise { const path = url.pathname; if (method === "GET" && (path === "/health" || path === "/health/live" || path === "/health/readiness")) { const [database, runnerDispatch, kafkaEventOutbox] = await Promise.all([store.health(), store.runnerDispatchStatus(), store.kafkaEventOutboxStatus()]); @@ -849,8 +907,9 @@ async function route({ method, url, body, store, sourceCommit, authSummary, runn const afterSeq = integerQuery(url, "afterSeq", 0); const limit = integerQuery(url, "limit", 100); const run = await store.getRun(runId); + const traceContext = await validatedDiagnosticTraceContext({ store, context: diagnosticTraceContext, run }); const items = await store.listEvents(runId, afterSeq, limit); - void emitAgentRunOtelSpan("projection_sync", run, process.env, { startTimeMs: startedAt, kind: 2, attributes: { "http.method": "GET", "http.route": "/api/v1/runs/:runId/events", "http.status_code": 200, afterSeq, limit, eventCount: items.length } }); + void emitAgentRunOtelSpan("projection_sync", run, process.env, { ...(traceContext ? { traceContext } : {}), startTimeMs: startedAt, kind: 2, attributes: { "http.method": "GET", "http.route": "/api/v1/runs/:runId/events", "http.status_code": 200, afterSeq, limit, eventCount: items.length } }); return { items: items as unknown as JsonValue }; } const runResultMatch = path.match(/^\/api\/v1\/runs\/([^/]+)\/result$/u); @@ -858,10 +917,20 @@ async function route({ method, url, body, store, sourceCommit, authSummary, runn const startedAt = Date.now(); const runId = runResultMatch[1] ?? ""; const run = await store.getRun(runId); - const commandId = url.searchParams.get("commandId") ?? undefined; + const commandIds = url.searchParams.getAll("commandId"); + const commandId = commandIds[0] ?? null; + if (diagnosticTraceContext !== null) { + if (commandIds.length !== 1 || commandId === null || commandId.length === 0) { + throw diagnosticTargetMismatch("diagnostic run result requires exactly one non-empty commandId query"); + } + if (commandId !== diagnosticTraceContext.targetCommandId) { + throw diagnosticTargetMismatch("diagnostic run result commandId does not match the diagnostic target command"); + } + } const command = commandId ? await store.getCommand(commandId) : null; - const result = await buildRunResult(store, runId, commandId) as JsonValue; - void emitAgentRunOtelSpan("command_result", run, process.env, { command, startTimeMs: startedAt, kind: 2, attributes: { "http.method": "GET", "http.route": "/api/v1/runs/:runId/result", "http.status_code": 200, ...commandResultOtelAttributes(result) } }); + const traceContext = await validatedDiagnosticTraceContext({ store, context: diagnosticTraceContext, run, command }); + const result = await buildRunResult(store, runId, commandId ?? undefined) as JsonValue; + void emitAgentRunOtelSpan("command_result", run, process.env, { command, ...(traceContext ? { traceContext } : {}), startTimeMs: startedAt, kind: 2, attributes: { "http.method": "GET", "http.route": "/api/v1/runs/:runId/result", "http.status_code": 200, ...commandResultOtelAttributes(result) } }); return result; } const runCancelMatch = path.match(/^\/api\/v1\/runs\/([^/]+)\/cancel$/u); @@ -953,8 +1022,9 @@ async function route({ method, url, body, store, sourceCommit, authSummary, runn const runId = commandResultMatch[1] ?? ""; const commandId = commandResultMatch[2] ?? ""; const [run, command] = await Promise.all([store.getRun(runId), store.getCommand(commandId)]); + const traceContext = await validatedDiagnosticTraceContext({ store, context: diagnosticTraceContext, run, command }); const result = await buildRunResult(store, runId, commandId) as JsonValue; - void emitAgentRunOtelSpan("command_result", run, process.env, { command, startTimeMs: startedAt, kind: 2, attributes: { "http.method": "GET", "http.route": "/api/v1/runs/:runId/commands/:commandId/result", "http.status_code": 200, ...commandResultOtelAttributes(result) } }); + void emitAgentRunOtelSpan("command_result", run, process.env, { command, ...(traceContext ? { traceContext } : {}), startTimeMs: startedAt, kind: 2, attributes: { "http.method": "GET", "http.route": "/api/v1/runs/:runId/commands/:commandId/result", "http.status_code": 200, ...commandResultOtelAttributes(result) } }); return result; } if (method === "POST" && path === "/api/v1/runners/register") return await store.registerRunner(asRecord(body ?? {}, "runner")) as unknown as JsonValue; diff --git a/src/selftest/cases/77-diagnostic-trace-context.ts b/src/selftest/cases/77-diagnostic-trace-context.ts new file mode 100644 index 0000000..2732520 --- /dev/null +++ b/src/selftest/cases/77-diagnostic-trace-context.ts @@ -0,0 +1,330 @@ +import assert from "node:assert/strict"; +import { createServer } from "node:http"; +import type { AddressInfo } from "node:net"; +import { + AGENTRUN_DIAGNOSTIC_OPERATION_HEADER, + AGENTRUN_DIAGNOSTIC_TARGET_BUSINESS_TRACE_ID_HEADER, + AGENTRUN_DIAGNOSTIC_TARGET_COMMAND_ID_HEADER, + AGENTRUN_TRACE_PURPOSE_HEADER, + agentRunOtelTraceContext, +} from "../../common/otel-trace.js"; +import type { JsonRecord, JsonValue } from "../../common/types.js"; +import { startManagerServer } from "../../mgr/server.js"; +import { MemoryAgentRunStore } from "../../mgr/store.js"; +import type { SelfTestCase } from "../harness.js"; + +const BUSINESS_TRACE_ID = "trc_diagnostic_context_selftest"; +const DIAGNOSTIC_OPERATION = "unidesk.diagnose-code-agent"; +const DIAGNOSTIC_TRACE_A = "11111111111111111111111111111111"; +const DIAGNOSTIC_TRACE_B = "22222222222222222222222222222222"; +const DIAGNOSTIC_PARENT_A = "aaaaaaaaaaaaaaaa"; +const DIAGNOSTIC_PARENT_B = "bbbbbbbbbbbbbbbb"; + +class TrackingMemoryAgentRunStore extends MemoryAgentRunStore { + createRunCalls = 0; + + override createRun(input: Parameters[0]) { + this.createRunCalls += 1; + return super.createRun(input); + } +} + +const selfTest: SelfTestCase = async () => { + const collector = await startOtlpCollector(); + const previousEndpoint = process.env.AGENTRUN_OTEL_EXPORTER_OTLP_TRACES_ENDPOINT; + process.env.AGENTRUN_OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = collector.endpoint; + const store = new TrackingMemoryAgentRunStore(); + const runInput: Parameters[0] = { + tenantId: "unidesk", + projectId: "pikasTech/agentrun", + workspaceRef: { kind: "host-path", path: "/tmp/agentrun-diagnostic-selftest" }, + providerId: "NC01", + backendProfile: "codex", + executionPolicy: { + sandbox: "workspace-write", + approval: "never", + timeoutMs: 1_000, + network: "none", + secretScope: { allowCredentialEcho: false, providerCredentials: [] }, + }, + traceSink: { kind: "hwlab", traceId: BUSINESS_TRACE_ID }, + }; + const run = store.createRun(runInput); + const command = store.createCommand(run.id, { + type: "turn", + payload: { prompt: "diagnostic trace context", traceId: BUSINESS_TRACE_ID }, + idempotencyKey: "diagnostic-trace-context", + }); + store.appendEvent(run.id, "assistant_message", { commandId: command.id, text: "diagnostic fixture response" }); + store.finishCommand(command.id, { terminalStatus: "completed", failureKind: null, failureMessage: null, threadId: "thread_diagnostic", turnId: "turn_diagnostic" }); + const manager = await startManagerServer({ + port: 0, + host: "127.0.0.1", + sourceCommit: "self-test", + store, + runnerJobDefaults: { + namespace: "agentrun-v02", + managerUrl: "http://agentrun-mgr.agentrun-v02.svc.cluster.local:8080", + runnerApiKeySecretRef: { name: "agentrun-v02-api-key", key: "HWLAB_API_KEY" }, + image: "registry.invalid/agentrun:self-test", + serviceAccountName: "agentrun-v02-runner", + jobNamePrefix: "agentrun-v02-runner", + lane: "v0.2", + }, + }); + + try { + const paths = [ + `/api/v1/runs/${run.id}/commands/${command.id}/result`, + `/api/v1/runs/${run.id}/result?commandId=${command.id}`, + `/api/v1/runs/${run.id}/events?afterSeq=0&limit=100`, + ]; + const ordinaryCommandDetail = asRecord(await getData(manager.baseUrl, `/api/v1/runs/${run.id}/commands/${command.id}`)); + assert.equal(ordinaryCommandDetail.id, command.id); + const ordinary = await Promise.all(paths.map((path) => getData(manager.baseUrl, path))); + await waitForSpanCount(collector.payloads, 3); + const businessContext = agentRunOtelTraceContext(run, command); + const initialBusinessSpans = spansForTrace(collector.payloads, businessContext.traceId); + assert.equal(initialBusinessSpans.length, 3); + assert.ok(initialBusinessSpans.every((span) => span.parentSpanId === businessContext.parentSpanId)); + assert.ok(initialBusinessSpans.every((span) => span.links === undefined)); + assert.ok(initialBusinessSpans.every((span) => spanAttributes(span).traceId === BUSINESS_TRACE_ID)); + + const diagnosticA = diagnosticHeaders(command.id, DIAGNOSTIC_TRACE_A, DIAGNOSTIC_PARENT_A, BUSINESS_TRACE_ID); + const diagnosticB = diagnosticHeaders(command.id, DIAGNOSTIC_TRACE_B, DIAGNOSTIC_PARENT_B, BUSINESS_TRACE_ID); + const firstDiagnostic = await Promise.all(paths.map((path) => getData(manager.baseUrl, path, diagnosticA))); + const secondDiagnostic = await Promise.all(paths.map((path) => getData(manager.baseUrl, path, diagnosticB))); + await waitForSpanCount(collector.payloads, 9); + + assert.deepEqual(authoritySummary(firstDiagnostic), authoritySummary(secondDiagnostic)); + assert.deepEqual(authoritySummary(firstDiagnostic), authoritySummary(ordinary)); + assert.equal(spansForTrace(collector.payloads, businessContext.traceId).length, 3); + assertDiagnosticSpans(spansForTrace(collector.payloads, DIAGNOSTIC_TRACE_A), DIAGNOSTIC_PARENT_A, businessContext, command.id); + assertDiagnosticSpans(spansForTrace(collector.payloads, DIAGNOSTIC_TRACE_B), DIAGNOSTIC_PARENT_B, businessContext, command.id); + + const latestCommand = store.createCommand(run.id, { + type: "turn", + payload: { prompt: "diagnostic latest-command guard", traceId: BUSINESS_TRACE_ID }, + idempotencyKey: "diagnostic-latest-command-guard", + }); + store.finishCommand(latestCommand.id, { terminalStatus: "failed", failureKind: "backend-timeout", failureMessage: "latest command must not replace the diagnostic target" }); + const spanCountBeforePinnedRunResult = allSpans(collector.payloads).length; + const pinnedRunResult = asRecord(await getData(manager.baseUrl, `/api/v1/runs/${run.id}/result?commandId=${command.id}`, diagnosticA)); + assert.equal(pinnedRunResult.commandId, command.id); + assert.equal(pinnedRunResult.terminalStatus, "completed"); + assert.notEqual(pinnedRunResult.commandId, latestCommand.id); + await waitForSpanCount(collector.payloads, spanCountBeforePinnedRunResult + 1); + + const spanCountBeforeRejectedReads = allSpans(collector.payloads).length; + const wrongBusiness = diagnosticHeaders(command.id, "33333333333333333333333333333333", "cccccccccccccccc", "trc_wrong_target"); + const wrongCommand = diagnosticHeaders("cmd_wrong_target", "44444444444444444444444444444444", "dddddddddddddddd", BUSINESS_TRACE_ID); + const writeDiagnostic = diagnosticHeaders(command.id, "55555555555555555555555555555555", "eeeeeeeeeeeeeeee", BUSINESS_TRACE_ID); + const createRunCallsBeforeRejectedWrites = store.createRunCalls; + await assertRejected(manager.baseUrl, paths[0] ?? "", wrongBusiness, "diagnostic-trace-target-mismatch"); + await assertRejected(manager.baseUrl, paths[0] ?? "", wrongCommand, "diagnostic-trace-target-mismatch"); + await assertRejected(manager.baseUrl, `/api/v1/runs/${run.id}/commands/${command.id}`, diagnosticA, "diagnostic-trace-route-unsupported"); + await assertRejected(manager.baseUrl, `/api/v1/runs/${run.id}/result`, diagnosticA, "diagnostic-trace-target-mismatch"); + await assertRejected(manager.baseUrl, `/api/v1/runs/${run.id}/result?commandId=`, diagnosticA, "diagnostic-trace-target-mismatch"); + await assertRejected(manager.baseUrl, `/api/v1/runs/${run.id}/result?commandId=${command.id}&commandId=${command.id}`, diagnosticA, "diagnostic-trace-target-mismatch"); + await assertRejected(manager.baseUrl, `/api/v1/runs/${run.id}/result?commandId=${latestCommand.id}`, diagnosticA, "diagnostic-trace-target-mismatch"); + await assertRejected(manager.baseUrl, "/api/v1/runs", writeDiagnostic, "diagnostic-trace-context-write-denied", "POST", runInput); + await waitForSpanCount(collector.payloads, spanCountBeforeRejectedReads + 8); + assert.equal(spansForTrace(collector.payloads, businessContext.traceId).length, 3); + const rejectedSpans = allSpans(collector.payloads).filter((span) => span.name === "diagnostic_request_rejected"); + assert.equal(rejectedSpans.length, 8); + for (const span of rejectedSpans) { + assert.equal((asRecord(span.status).code), 2); + assert.equal(Array.isArray(span.links), true); + assert.equal(spanAttributes(span).failureKind, "schema-invalid"); + assert.match(String(spanAttributes(span)["diagnostic.rejection_reason"]), /^diagnostic-trace-/u); + } + const emptyDiagnostic = diagnosticHeadersWithValue(diagnosticA, ""); + const whitespaceDiagnostic = diagnosticHeadersWithValue(diagnosticA, " "); + const repeatedDiagnostic = new Headers(diagnosticA); + repeatedDiagnostic.append(AGENTRUN_TRACE_PURPOSE_HEADER, "diagnostic"); + await assertRejected(manager.baseUrl, "/api/v1/runs", emptyDiagnostic, "diagnostic-trace-context-invalid", "POST", runInput); + await assertRejected(manager.baseUrl, "/api/v1/runs", whitespaceDiagnostic, "diagnostic-trace-context-invalid", "POST", runInput); + await assertRejected(manager.baseUrl, "/api/v1/runs", repeatedDiagnostic, "diagnostic-trace-context-invalid", "POST", runInput); + assert.equal(store.createRunCalls, createRunCallsBeforeRejectedWrites); + const malformedTraceparent = { ...diagnosticA, traceparent: "00-00000000000000000000000000000000-0000000000000000-01" }; + await assertRejected(manager.baseUrl, paths[0] ?? "", malformedTraceparent, "diagnostic-trace-context-invalid"); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(allSpans(collector.payloads).length, spanCountBeforeRejectedReads + 8); + + return { + name: "diagnostic-trace-context", + tests: [ + "ordinary-read-business-trace-preserved", + "diagnostic-read-independent-trace", + "diagnostic-span-standard-link", + "diagnostic-result-events-authority-stable", + "diagnostic-target-mismatch-fail-closed", + "diagnostic-target-command-mismatch-fail-closed", + "diagnostic-command-detail-remains-ordinary", + "diagnostic-run-result-command-query-fail-closed", + "diagnostic-write-request-fail-closed", + "diagnostic-empty-and-repeated-headers-fail-closed", + "diagnostic-rejection-error-span", + "malformed-diagnostic-context-no-span", + ], + }; + } finally { + await new Promise((resolve) => manager.server.close(() => resolve())); + await collector.close(); + if (previousEndpoint === undefined) delete process.env.AGENTRUN_OTEL_EXPORTER_OTLP_TRACES_ENDPOINT; + else process.env.AGENTRUN_OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = previousEndpoint; + } +}; + +function diagnosticHeaders(commandId: string, traceId: string, parentSpanId: string, targetBusinessTraceId: string): Record { + return { + traceparent: `00-${traceId}-${parentSpanId}-01`, + [AGENTRUN_TRACE_PURPOSE_HEADER]: "diagnostic", + [AGENTRUN_DIAGNOSTIC_OPERATION_HEADER]: DIAGNOSTIC_OPERATION, + [AGENTRUN_DIAGNOSTIC_TARGET_BUSINESS_TRACE_ID_HEADER]: targetBusinessTraceId, + [AGENTRUN_DIAGNOSTIC_TARGET_COMMAND_ID_HEADER]: commandId, + }; +} + +function diagnosticHeadersWithValue(base: Record, value: string): Record { + return { + traceparent: base.traceparent ?? "", + [AGENTRUN_TRACE_PURPOSE_HEADER]: value, + [AGENTRUN_DIAGNOSTIC_OPERATION_HEADER]: value, + [AGENTRUN_DIAGNOSTIC_TARGET_BUSINESS_TRACE_ID_HEADER]: value, + [AGENTRUN_DIAGNOSTIC_TARGET_COMMAND_ID_HEADER]: value, + }; +} + +async function getData(baseUrl: string, path: string, headers: Record = {}): Promise { + const response = await fetch(new URL(path, baseUrl), { headers }); + const envelope = await response.json() as JsonRecord; + assert.equal(response.status, 200, JSON.stringify(envelope)); + assert.equal(envelope.ok, true, JSON.stringify(envelope)); + return envelope.data ?? null; +} + +async function assertRejected(baseUrl: string, path: string, headers: Record | Headers, expectedReason: string, method = "GET", body?: unknown): Promise { + const requestHeaders = new Headers(headers); + if (body !== undefined) requestHeaders.set("content-type", "application/json"); + const response = await fetch(new URL(path, baseUrl), { + method, + headers: requestHeaders, + ...(body === undefined ? {} : { body: JSON.stringify(body) }), + }); + const envelope = await response.json() as JsonRecord; + assert.equal(response.status, 400, JSON.stringify(envelope)); + assert.equal(envelope.ok, false); + assert.equal(envelope.failureKind, "schema-invalid"); + assert.equal((((envelope.error as JsonRecord).details as JsonRecord).reason), expectedReason); +} + +function authoritySummary(values: JsonValue[]): JsonValue[] { + return values.map((value, index) => { + const record = asRecord(value); + if (index === 2) { + const items = Array.isArray(record.items) ? record.items.map((item) => { + const event = asRecord(item); + return { seq: event.seq ?? null, type: event.type ?? null }; + }) : []; + return { items }; + } + return { + commandId: record.commandId ?? null, + terminalStatus: record.terminalStatus ?? null, + failureKind: record.failureKind ?? null, + eventCount: record.eventCount ?? null, + lastSeq: record.lastSeq ?? null, + }; + }); +} + +function assertDiagnosticSpans(spans: JsonRecord[], expectedParentSpanId: string, businessContext: ReturnType, commandId: string): void { + assert.equal(spans.length, 3); + assert.deepEqual(spans.map((span) => span.name).sort(), ["command_result", "command_result", "projection_sync"]); + for (const span of spans) { + assert.equal(span.parentSpanId, expectedParentSpanId); + const attributes = spanAttributes(span); + assert.equal(Object.hasOwn(attributes, "traceId"), false); + assert.equal(attributes["otel.trace_id"], span.traceId); + assert.equal(attributes["diagnostic.target_business_trace_id"], BUSINESS_TRACE_ID); + assert.equal(attributes["diagnostic.target_otel_trace_id"], businessContext.traceId); + assert.equal(attributes["diagnostic.operation"], DIAGNOSTIC_OPERATION); + assert.equal(attributes.commandId, commandId); + const links = Array.isArray(span.links) ? span.links.map(asRecord) : []; + assert.equal(links.length, 1); + assert.equal(links[0]?.traceId, businessContext.traceId); + assert.equal(links[0]?.spanId, businessContext.parentSpanId); + assert.match(String(links[0]?.traceId), /^[0-9a-f]{32}$/u); + assert.match(String(links[0]?.spanId), /^[0-9a-f]{16}$/u); + const linkAttributes = otlpAttributes(links[0]?.attributes); + assert.equal(linkAttributes["diagnostic.target_business_trace_id"], BUSINESS_TRACE_ID); + assert.equal(linkAttributes["diagnostic.target_otel_trace_id"], businessContext.traceId); + assert.equal(linkAttributes["diagnostic.operation"], DIAGNOSTIC_OPERATION); + assert.equal(linkAttributes.commandId, commandId); + } +} + +function spanAttributes(span: JsonRecord): Record { + return otlpAttributes(span.attributes); +} + +function otlpAttributes(value: JsonValue | undefined): Record { + const output: Record = {}; + if (!Array.isArray(value)) return output; + for (const item of value) { + const attribute = asRecord(item); + const key = typeof attribute.key === "string" ? attribute.key : null; + const wrapped = asRecord(attribute.value); + if (!key) continue; + if (typeof wrapped.stringValue === "string") output[key] = wrapped.stringValue; + else if (typeof wrapped.intValue === "string") output[key] = wrapped.intValue; + else if (typeof wrapped.doubleValue === "number") output[key] = wrapped.doubleValue; + else if (typeof wrapped.boolValue === "boolean") output[key] = wrapped.boolValue; + } + return output; +} + +function allSpans(payloads: JsonRecord[]): JsonRecord[] { + return payloads.flatMap((payload) => { + const resourceSpans = Array.isArray(payload.resourceSpans) ? payload.resourceSpans.map(asRecord) : []; + return resourceSpans.flatMap((resourceSpan) => { + const scopeSpans = Array.isArray(resourceSpan.scopeSpans) ? resourceSpan.scopeSpans.map(asRecord) : []; + return scopeSpans.flatMap((scopeSpan) => Array.isArray(scopeSpan.spans) ? scopeSpan.spans.map(asRecord) : []); + }); + }); +} + +function spansForTrace(payloads: JsonRecord[], traceId: string): JsonRecord[] { + return allSpans(payloads).filter((span) => span.traceId === traceId); +} + +async function waitForSpanCount(payloads: JsonRecord[], expected: number): Promise { + const deadline = Date.now() + 3_000; + while (allSpans(payloads).length < expected && Date.now() < deadline) await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(allSpans(payloads).length, expected); +} + +async function startOtlpCollector(): Promise<{ endpoint: string; payloads: JsonRecord[]; close: () => Promise }> { + const payloads: JsonRecord[] = []; + const server = createServer(async (request, response) => { + const chunks: Buffer[] = []; + for await (const chunk of request) chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); + payloads.push(JSON.parse(Buffer.concat(chunks).toString("utf8")) as JsonRecord); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address() as AddressInfo; + return { + endpoint: `http://${address.address}:${address.port}/v1/traces`, + payloads, + close: async () => await new Promise((resolve) => server.close(() => resolve())), + }; +} + +function asRecord(value: unknown): JsonRecord { + return typeof value === "object" && value !== null && !Array.isArray(value) ? value as JsonRecord : {}; +} + +export default selfTest;