@senad-d/observme 0.1.4 → 0.1.5

package/src/pi/session-correlation.ts

@@ -0,0 +1,167 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { OBSERVME_CORRELATION_ENTRY_TYPE } from "../semconv/values.ts";
+import type { AgentLineageContext } from "./agent-lineage.ts";
+import type { MinimalSessionCorrelation } from "./handler-types.ts";
+
+export { OBSERVME_CORRELATION_ENTRY_TYPE } from "../semconv/values.ts";
+export const OBSERVME_CORRELATION_ENTRY_SCHEMA_VERSION = 1;
+
+const maximumCorrelationValueLength = 128;
+const correlationValuePattern = /^[A-Za-z0-9._:-]+$/u;
+const correlationDataKeys = new Set([
+  "schemaVersion",
+  "workflowId",
+  "agentId",
+  "parentAgentId",
+  "rootAgentId",
+  "parentSessionId",
+  "depth",
+  "spawnId",
+  "capability",
+]);
+
+export interface ObservMeCorrelationEntryData extends MinimalSessionCorrelation {
+  readonly schemaVersion: typeof OBSERVME_CORRELATION_ENTRY_SCHEMA_VERSION;
+  readonly workflowId: string;
+  readonly agentId: string;
+  readonly rootAgentId: string;
+  readonly depth: number;
+}
+
+export function readLatestSessionCorrelation(entries: readonly unknown[] | undefined): MinimalSessionCorrelation | undefined {
+  if (!entries) return undefined;
+
+  for (let index = entries.length - 1; index >= 0; index -= 1) {
+    const correlation = readSessionCorrelationEntry(entries[index]);
+    if (correlation) return correlation;
+  }
+
+  return undefined;
+}
+
+export function appendSessionCorrelationEntry(
+  appendEntry: ExtensionAPI["appendEntry"] | undefined,
+  lineage: AgentLineageContext,
+  recovered: MinimalSessionCorrelation | undefined,
+): void {
+  if (!appendEntry) return;
+
+  const data = buildSessionCorrelationEntryData(lineage);
+  if (!data || correlationsEqual(data, recovered)) return;
+
+  try {
+    appendEntry(OBSERVME_CORRELATION_ENTRY_TYPE, data);
+  } catch {
+    return;
+  }
+}
+
+export function buildSessionCorrelationEntryData(
+  lineage: AgentLineageContext,
+): ObservMeCorrelationEntryData | undefined {
+  return normalizeSessionCorrelationData({
+    schemaVersion: OBSERVME_CORRELATION_ENTRY_SCHEMA_VERSION,
+    workflowId: lineage.workflowId,
+    agentId: lineage.agentId,
+    parentAgentId: lineage.parentAgentId,
+    rootAgentId: lineage.rootAgentId,
+    parentSessionId: lineage.parentSessionId,
+    depth: lineage.depth,
+    spawnId: lineage.spawnId,
+    capability: lineage.capability,
+  });
+}
+
+function readSessionCorrelationEntry(value: unknown): ObservMeCorrelationEntryData | undefined {
+  if (!isRecord(value)) return undefined;
+  if (value.type !== "custom" || value.customType !== OBSERVME_CORRELATION_ENTRY_TYPE) return undefined;
+  return normalizeSessionCorrelationData(value.data);
+}
+
+function normalizeSessionCorrelationData(value: unknown): ObservMeCorrelationEntryData | undefined {
+  if (!isRecord(value) || !hasOnlyCorrelationDataKeys(value)) return undefined;
+  if (value.schemaVersion !== OBSERVME_CORRELATION_ENTRY_SCHEMA_VERSION) return undefined;
+
+  const workflowId = readCorrelationString(value, "workflowId");
+  const agentId = readCorrelationString(value, "agentId");
+  const parentAgentId = readOptionalCorrelationString(value, "parentAgentId");
+  const rootAgentId = readCorrelationString(value, "rootAgentId");
+  const parentSessionId = readOptionalCorrelationString(value, "parentSessionId");
+  const depth = readCorrelationDepth(value.depth);
+  const spawnId = readOptionalCorrelationString(value, "spawnId");
+  const capability = readOptionalCorrelationString(value, "capability");
+
+  if (!workflowId || !agentId || !rootAgentId || depth === undefined) return undefined;
+  if (parentAgentId === null || parentSessionId === null || spawnId === null || capability === null) return undefined;
+  if (!isConsistentCorrelationLineage(agentId, parentAgentId, rootAgentId, depth)) return undefined;
+
+  return withoutUndefinedValues({
+    schemaVersion: OBSERVME_CORRELATION_ENTRY_SCHEMA_VERSION,
+    workflowId,
+    agentId,
+    parentAgentId,
+    rootAgentId,
+    parentSessionId,
+    depth,
+    spawnId,
+    capability,
+  }) as ObservMeCorrelationEntryData;
+}
+
+function hasOnlyCorrelationDataKeys(value: Record<string, unknown>): boolean {
+  return Object.keys(value).every(key => correlationDataKeys.has(key));
+}
+
+function readCorrelationString(value: Record<string, unknown>, key: string): string | undefined {
+  const candidate = value[key];
+  if (typeof candidate !== "string") return undefined;
+  if (candidate.length === 0 || candidate.length > maximumCorrelationValueLength) return undefined;
+  return correlationValuePattern.test(candidate) ? candidate : undefined;
+}
+
+function readOptionalCorrelationString(
+  value: Record<string, unknown>,
+  key: string,
+): string | undefined | null {
+  if (!Object.hasOwn(value, key) || value[key] === undefined) return undefined;
+  return readCorrelationString(value, key) ?? null;
+}
+
+function readCorrelationDepth(value: unknown): number | undefined {
+  if (typeof value !== "number" || !Number.isInteger(value)) return undefined;
+  return value >= 0 && value <= 64 ? value : undefined;
+}
+
+function isConsistentCorrelationLineage(
+  agentId: string,
+  parentAgentId: string | undefined,
+  rootAgentId: string,
+  depth: number,
+): boolean {
+  if (depth === 0) return parentAgentId === undefined && rootAgentId === agentId;
+  return parentAgentId !== undefined && parentAgentId !== agentId;
+}
+
+function correlationsEqual(
+  left: ObservMeCorrelationEntryData,
+  right: MinimalSessionCorrelation | undefined,
+): boolean {
+  if (!right) return false;
+
+  return left.workflowId === right.workflowId
+    && left.agentId === right.agentId
+    && left.parentAgentId === right.parentAgentId
+    && left.rootAgentId === right.rootAgentId
+    && left.parentSessionId === right.parentSessionId
+    && left.depth === right.depth
+    && left.spawnId === right.spawnId
+    && left.capability === right.capability;
+}
+
+function withoutUndefinedValues<T extends Record<string, unknown>>(value: T): T {
+  return Object.fromEntries(Object.entries(value).filter(entry => entry[1] !== undefined)) as T;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}

package/src/pi/subagent-spawn.ts

@@ -24,7 +24,7 @@ import {
 } from "../semconv/values.ts";
 import { BoundedMap } from "../util/bounded-map.ts";
 import type { AgentChildStatus, AgentTreeNode, AgentTreeSummary } from "./agent-tree-tracker.ts";
-import { AgentTreeTracker } from "./agent-tree-tracker.ts";
+import { AgentTreeTracker, isAgentStatusTransitionAllowed } from "./agent-tree-tracker.ts";
 import type { AgentLineageContext, AgentRole } from "./agent-lineage.ts";
 import { createAgentLineageContext, createPropagationEnvironment, sanitizePropagationEnvironment } from "./agent-lineage.ts";
 import { recordActiveSpanEnd, recordActiveSpanStart } from "./handler-internals.ts";
@@ -44,6 +44,7 @@ export type {
 
 export type SubagentSpawnType = "command" | "tool" | "extension" | "unknown";
 export type AgentJoinStatus = "completed" | "failed" | "cancelled" | "timeout" | "unknown" | "waiting";
+export type AgentTerminalStatus = Extract<AgentChildStatus, "completed" | "failed" | "cancelled">;
 export type AttributePrimitive = boolean | number | string | string[];
 export type AttributeMap = Record<string, AttributePrimitive>;
 
@@ -125,13 +126,35 @@ export interface StartedSubagentSpawn {
   readonly attributes: AttributeMap;
 }
 
+export interface SubagentSpawnIdentity {
+  readonly spawnId: string;
+  readonly childAgentId: string;
+}
+
 export interface CompleteSubagentSpawnOptions {
   readonly childAgentId?: string;
-  readonly childStatus?: AgentChildStatus;
-  readonly outcome?: "completed" | "failed" | "cancelled";
+  readonly childStatus?: AgentTerminalStatus;
+  readonly outcome?: AgentTerminalStatus;
   readonly now?: () => number;
 }
 
+export type SubagentTransitionFailureReason =
+  | "spawn_not_found"
+  | "child_agent_mismatch"
+  | "invalid_terminal_transition";
+
+export type SubagentTransitionResult =
+  | { readonly ok: true }
+  | { readonly ok: false; readonly reason: SubagentTransitionFailureReason };
+
+interface ResolvedCompletionTransition {
+  readonly ok: true;
+  readonly childAgentId: string;
+  readonly status: AgentTerminalStatus;
+}
+
+type CompletionTransitionResult = ResolvedCompletionTransition | Extract<SubagentTransitionResult, { readonly ok: false }>;
+
 export interface FailSubagentSpawnOptions {
   readonly childAgentId?: string;
   readonly errorClass?: string;
@@ -219,12 +242,18 @@ export async function runSubagentWithObservability<Result>(
   }
 }
 
+export function resolveSubagentSpawnIdentity(
+  options: Pick<StartSubagentSpawnOptions, "spawnId" | "childAgentId"> = {},
+): SubagentSpawnIdentity {
+  const spawnId = options.spawnId ?? `spawn-${randomUUID()}`;
+  return { spawnId, childAgentId: options.childAgentId ?? `child-${spawnId}` };
+}
+
 export function startSubagentSpawn(
   session: SubagentTelemetrySession,
   options: StartSubagentSpawnOptions = {},
 ): StartedSubagentSpawn {
-  const spawnId = options.spawnId ?? `spawn-${randomUUID()}`;
-  const childAgentId = options.childAgentId ?? `child-${spawnId}`;
+  const { spawnId, childAgentId } = resolveSubagentSpawnIdentity(options);
   const spawnReason = normalizeSpawnReason(options.spawnReason);
   const labels = subagentSpawnMetricLabels(session, options, spawnReason);
   const parentSpan = resolveSubagentParentSpan(session);
@@ -269,48 +298,55 @@ export function completeSubagentSpawn(
   session: SubagentTelemetrySession,
   spawnId: string,
   options: CompleteSubagentSpawnOptions = {},
-): void {
+): SubagentTransitionResult {
   const state = session.spans.activeSubagentSpawns.get(spawnId);
-  if (!state) return;
+  if (!state) return subagentTransitionFailure("spawn_not_found");
 
-  const childAgentId = options.childAgentId ?? state.childAgentId;
-  const childStatus = options.childStatus ?? "completed";
+  const transition = resolveCompletionTransition(session, state, options);
+  if (!transition.ok) return transition;
+
+  updateChildStatus(session, transition.childAgentId, transition.status);
+  recordChildFailureCompletion(session, transition.childAgentId, transition.status);
+  recordObsAgentsTreeState(session);
   const attributes = buildSubagentCompletionAttributes(
     session,
     spawnId,
-    childAgentId,
-    options.outcome ?? "completed",
+    transition.childAgentId,
+    transition.status,
     state.spawnReason,
   );
+  const eventName = terminalSpawnEventName(transition.status);
+  const severityText = transition.status === "completed" ? "INFO" : "ERROR";
 
-  updateChildStatus(session, childAgentId, childStatus);
-  recordChildFailureCompletion(session, childAgentId, childStatus);
-  recordObsAgentsTreeState(session);
   recordSubagentSpawnDuration(session, state, options);
   state.span.setAttributes(attributes);
-  state.span.setStatus({ code: SpanStatusCode.OK });
-  state.span.addEvent(LOG_EVENT_NAMES.AGENT_SPAWN_COMPLETED, attributes);
+  setTerminalSpawnStatus(state.span, transition.status);
+  state.span.addEvent(eventName, attributes);
   endSubagentSpan(session, state.span);
   session.spans.activeSubagentSpawns.delete(spawnId);
-  emitSubagentLog(session, LOG_EVENT_NAMES.AGENT_SPAWN_COMPLETED, attributes);
+  emitSubagentLog(session, eventName, attributes, severityText);
+  return subagentTransitionSuccess();
 }
 
 export function failSubagentSpawn(
   session: SubagentTelemetrySession,
   spawnId: string,
   options: FailSubagentSpawnOptions = {},
-): void {
+): SubagentTransitionResult {
   const state = session.spans.activeSubagentSpawns.get(spawnId);
-  if (!state) return;
+  if (!state) return subagentTransitionFailure("spawn_not_found");
 
   const childAgentId = options.childAgentId ?? state.childAgentId;
+  const transitionFailure = validateLauncherFailureTransition(session, state, childAgentId);
+  if (transitionFailure) return transitionFailure;
+
+  updateChildStatus(session, childAgentId, "failed");
+  recordObsAgentsTreeState(session);
   const attributes = {
     ...buildSubagentCompletionAttributes(session, spawnId, childAgentId, "failed", state.spawnReason),
     [LOG_ATTRIBUTES.ERROR_TYPE]: normalizeMetricLabel(options.errorClass ?? "subagent_spawn_error", "subagent_spawn_error"),
   };
 
-  updateChildStatus(session, childAgentId, "failed");
-  recordObsAgentsTreeState(session);
   recordSubagentSpawnDuration(session, state, options);
   state.span.setAttributes(attributes);
   state.span.setStatus({ code: SpanStatusCode.ERROR, message: String(attributes[LOG_ATTRIBUTES.ERROR_TYPE]) });
@@ -319,6 +355,7 @@ export function failSubagentSpawn(
   session.spans.activeSubagentSpawns.delete(spawnId);
   session.metrics.subagentSpawnFailures.add(1, subagentFailureMetricLabels(state.labels, attributes));
   emitSubagentLog(session, LOG_EVENT_NAMES.AGENT_SPAWN_FAILED, attributes, "ERROR");
+  return subagentTransitionSuccess();
 }
 
 export function buildSubagentPropagationEnvironment(
@@ -372,8 +409,8 @@ export function endAgentWait(
   session: SubagentTelemetrySession,
   waitId: string,
   options: AgentWaitJoinOptions = {},
-): void {
-  endWaitJoinSpan(session, waitId, options, "wait");
+): SubagentTransitionResult {
+  return endWaitJoinSpan(session, waitId, options, "wait");
 }
 
 export function recordAgentWait(session: SubagentTelemetrySession, options: AgentWaitJoinOptions = {}): StartedAgentWaitJoin {
@@ -393,8 +430,8 @@ export function endAgentJoin(
   session: SubagentTelemetrySession,
   joinId: string,
   options: AgentWaitJoinOptions = {},
-): void {
-  endWaitJoinSpan(session, joinId, options, "join");
+): SubagentTransitionResult {
+  return endWaitJoinSpan(session, joinId, options, "join");
 }
 
 export function recordAgentJoin(session: SubagentTelemetrySession, options: AgentWaitJoinOptions = {}): StartedAgentWaitJoin {
@@ -468,11 +505,13 @@ function endWaitJoinSpan(
   id: string,
   options: AgentWaitJoinOptions,
   kind: "wait" | "join",
-): void {
+): SubagentTransitionResult {
   const registry = waitJoinRegistry(session, kind);
   const state = registry.get(id);
-  if (!state) return;
+  if (!state) return subagentTransitionFailure("spawn_not_found");
 
+  const transition = validateWaitJoinTransition(session, options, kind);
+  if (!transition.ok) return transition;
   if (kind === "join" && options.childAgentId && options.childStatus) updateChildStatus(session, options.childAgentId, options.childStatus);
 
   const attributes = buildWaitJoinAttributes(session, options, state.reason);
@@ -489,6 +528,28 @@ function endWaitJoinSpan(
   recordWaitJoinDuration(session, durationMs, state.labels, kind);
   if (kind === "join") recordChildJoinAccounting(session, options);
   emitSubagentLog(session, eventName, attributes, waitJoinFailed(options) ? "ERROR" : "INFO");
+  return subagentTransitionSuccess();
+}
+
+function validateWaitJoinTransition(
+  session: SubagentTelemetrySession,
+  options: AgentWaitJoinOptions,
+  kind: "wait" | "join",
+): SubagentTransitionResult {
+  if (kind === "wait" || !options.childAgentId || !options.childStatus) return subagentTransitionSuccess();
+
+  if (options.spawnId) {
+    const spawn = session.spans.activeSubagentSpawns.get(options.spawnId);
+    if (spawn && spawn.childAgentId !== options.childAgentId) {
+      return subagentTransitionFailure("child_agent_mismatch");
+    }
+  }
+
+  const child = ensureAgentTree(session).getAgent(options.childAgentId);
+  if (child && !isAgentStatusTransitionAllowed(child.status, options.childStatus)) {
+    return subagentTransitionFailure("invalid_terminal_transition");
+  }
+  return subagentTransitionSuccess();
 }
 
 function waitJoinRegistry(
@@ -523,7 +584,7 @@ function recordSubagentSpawnDuration(
 function recordChildFailureCompletion(
   session: SubagentTelemetrySession,
   childAgentId: string,
-  childStatus: AgentChildStatus,
+  childStatus: AgentTerminalStatus,
 ): void {
   if (childStatus !== "failed") return;
   recordChildFailureAccounting(session, childAgentId, false);
@@ -644,11 +705,83 @@ function buildSubagentSpawnAttributes(
   });
 }
 
+function resolveCompletionTransition(
+  session: SubagentTelemetrySession,
+  state: SubagentSpawnState,
+  options: CompleteSubagentSpawnOptions,
+): CompletionTransitionResult {
+  const childAgentId = options.childAgentId ?? state.childAgentId;
+  const requestedStatus: unknown = options.childStatus ?? options.outcome ?? "completed";
+  if (!isAgentTerminalStatus(requestedStatus)) return subagentTransitionFailure("invalid_terminal_transition");
+  if (options.childStatus !== undefined && options.outcome !== undefined && options.childStatus !== options.outcome) {
+    return subagentTransitionFailure("invalid_terminal_transition");
+  }
+
+  const transitionFailure = validateSpawnChildTransition(session, state, childAgentId, requestedStatus);
+  if (transitionFailure) return transitionFailure;
+  return { ok: true, childAgentId, status: requestedStatus };
+}
+
+function validateSpawnChildTransition(
+  session: SubagentTelemetrySession,
+  state: SubagentSpawnState,
+  childAgentId: string,
+  status: AgentTerminalStatus,
+): Extract<SubagentTransitionResult, { readonly ok: false }> | undefined {
+  if (childAgentId !== state.childAgentId) return subagentTransitionFailure("child_agent_mismatch");
+
+  const child = ensureAgentTree(session).getAgent(childAgentId);
+  if (!child || !isAgentStatusTransitionAllowed(child.status, status)) {
+    return subagentTransitionFailure("invalid_terminal_transition");
+  }
+  return undefined;
+}
+
+function validateLauncherFailureTransition(
+  session: SubagentTelemetrySession,
+  state: SubagentSpawnState,
+  childAgentId: string,
+): Extract<SubagentTransitionResult, { readonly ok: false }> | undefined {
+  if (childAgentId !== state.childAgentId) return subagentTransitionFailure("child_agent_mismatch");
+
+  const child = ensureAgentTree(session).getAgent(childAgentId);
+  if (child?.status !== "starting") return subagentTransitionFailure("invalid_terminal_transition");
+  return undefined;
+}
+
+function isAgentTerminalStatus(value: unknown): value is AgentTerminalStatus {
+  return value === "completed" || value === "failed" || value === "cancelled";
+}
+
+function terminalSpawnEventName(status: AgentTerminalStatus): string {
+  if (status === "completed") return LOG_EVENT_NAMES.AGENT_SPAWN_COMPLETED;
+  if (status === "failed") return LOG_EVENT_NAMES.AGENT_SPAWN_FAILED;
+  return LOG_EVENT_NAMES.AGENT_SPAWN_CANCELLED;
+}
+
+function setTerminalSpawnStatus(span: Span, status: AgentTerminalStatus): void {
+  if (status === "completed") {
+    span.setStatus({ code: SpanStatusCode.OK });
+    return;
+  }
+  span.setStatus({ code: SpanStatusCode.ERROR, message: status });
+}
+
+function subagentTransitionSuccess(): SubagentTransitionResult {
+  return { ok: true };
+}
+
+function subagentTransitionFailure(
+  reason: SubagentTransitionFailureReason,
+): Extract<SubagentTransitionResult, { readonly ok: false }> {
+  return { ok: false, reason };
+}
+
 function buildSubagentCompletionAttributes(
   session: SubagentTelemetrySession,
   spawnId: string,
   childAgentId: string,
-  outcome: "completed" | "failed" | "cancelled",
+  outcome: AgentTerminalStatus,
   spawnReason: SubagentSpawnReason,
 ): AttributeMap {
   const summary = ensureAgentTree(session).summarize(session.lineage.rootAgentId);
@@ -736,8 +869,8 @@ function updateChildStatus(
   session: SubagentTelemetrySession,
   childAgentId: string,
   status: AgentChildStatus,
-): void {
-  ensureAgentTree(session).updateStatus(childAgentId, status);
+): AgentTreeNode | undefined {
+  return ensureAgentTree(session).updateStatus(childAgentId, status);
 }
 
 function recordOrphanAgent(session: SubagentTelemetrySession, node: AgentTreeNode): void {