@senad-d/observme 0.1.4 → 0.1.5

package/src/pi/event-handlers/lifecycle.ts

@@ -1,5 +1,10 @@
 import { open } from "node:fs/promises";
 import { SpanStatusCode } from "@opentelemetry/api";
+import type {
+  ExtensionContext,
+  SessionShutdownEvent,
+  SessionStartEvent,
+} from "@earendil-works/pi-coding-agent";
 import {
   clearObsAgentsRuntimeState,
   startObsAgentsRuntimeState,
@@ -25,10 +30,10 @@ import {
 import type { ObservMeConfig } from "../../config/schema.ts";
 import { emitUnsafeCaptureWarning, normalizeConfigRejectionDiagnostic } from "../../config/validate.ts";
 import { EXTENSION_STATUS_KEY, EXTENSION_STATUS_VALUE } from "../../constants.ts";
-import type { BoundedOtelOperationResult } from "../../otel/shutdown.ts";
+import { ObservMeOtelStartupError } from "../../otel/sdk.ts";
+import type { BoundedOtelOperationResult, OtelOperationSettlement } from "../../otel/shutdown.ts";
 import {
   AGENT_LINEAGE_ATTRIBUTES,
-  COMMON_SPAN_ATTRIBUTES,
   CONFIG_ATTRIBUTES,
   LOG_ATTRIBUTES,
   SESSION_ATTRIBUTES,
@@ -55,7 +60,6 @@ import {
   readSpanId,
   readSpanTraceId,
   readString,
-  readUnknown,
   resolveModelId,
   resolveModelProvider,
   resolveSessionFilePath,
@@ -71,15 +75,19 @@ import {
   startSessionTelemetry,
   workflowFailed,
 } from "../handler-runtime.ts";
+import {
+  appendSessionCorrelationEntry,
+  readLatestSessionCorrelation,
+} from "../session-correlation.ts";
 import type { HandlerRegistrar, SerializedLifecycleQueue } from "../handler-runtime.ts";
 import type {
   AttributeMap,
-  Handler,
   HandlerSessionState,
   LoadSessionConfig,
   MinimalSessionCorrelation,
   ObservMeHandlerContext,
   ObservMeTelemetrySession,
+  PiHandler,
   RegisterHandlersOptions,
   SessionConfigLoadResult,
   SessionRecoveryHeader,
@@ -108,23 +116,33 @@ export function buildSessionAttributes(
   config: ObservMeConfig,
   lineage: ObservMeTelemetrySession["lineage"],
   recovery?: StartupRecoveryState,
+  initialThinkingLevel?: unknown,
 ): AttributeMap {
-  const cwd = recovery?.header?.cwd ?? readString(ctx, "cwd") ?? process.cwd();
-  const sessionId = recovery?.header?.id ?? resolveSessionId(event, ctx, lineage);
-  const parentSessionId = recovery?.header?.parentSession ?? readString(event, "parentSessionId") ?? lineage.parentSessionId;
-  const sessionFile = recovery?.sessionFile ?? resolveSessionFilePath(event, ctx);
+  const sessionManager = ctx.sessionManager;
+  const managerHeader = normalizeSessionHeader(sessionManager?.getHeader());
+  const header = managerHeader ?? recovery?.header;
+  const cwd = sessionManager?.getCwd() ?? header?.cwd ?? ctx.cwd ?? process.cwd();
+  const sessionId = sessionManager?.getSessionId() ?? header?.id ?? resolveSessionId(event, ctx, lineage);
+  const parentSessionId = header?.parentSession ?? lineage.parentSessionId;
+  const sessionFile = sessionManager?.getSessionFile() ?? recovery?.sessionFile ?? resolveSessionFilePath(event, ctx);
+  const sessionName = sessionManager
+    ? sessionManager.getSessionName() ?? "unknown"
+    : readString(event, "sessionName") ?? readString(event, "name") ?? "unknown";
+  const persisted = sessionManager
+    ? sessionFile !== undefined
+    : readBoolean(event, "persisted") ?? recovery?.resumed ?? false;
 
   return withoutUndefinedAttributes({
     [SESSION_ATTRIBUTES.PI_SESSION_ID]: sessionId,
-    [SESSION_ATTRIBUTES.PI_SESSION_NAME]: readString(event, "sessionName") ?? readString(event, "name") ?? "unknown",
+    [SESSION_ATTRIBUTES.PI_SESSION_NAME]: sessionName,
     [SESSION_ATTRIBUTES.PI_SESSION_CWD_HASH]: hashValue(cwd, config),
     [SESSION_ATTRIBUTES.PI_SESSION_PARENT_SESSION_HASH]: parentSessionId ? hashValue(parentSessionId, config) : "",
-    [SESSION_ATTRIBUTES.PI_SESSION_PERSISTED]: readBoolean(event, "persisted") ?? recovery?.resumed ?? false,
+    [SESSION_ATTRIBUTES.PI_SESSION_PERSISTED]: persisted,
     [SESSION_ATTRIBUTES.PI_SESSION_FILE_HASH]: sessionFile ? hashValue(sessionFile, config) : "",
-    [SESSION_ATTRIBUTES.PI_SESSION_VERSION]: readString(recovery?.header, "version") ?? readString(event, "sessionVersion") ?? readString(event, "version") ?? "unknown",
-    [SESSION_ATTRIBUTES.PI_MODEL_PROVIDER_CURRENT]: resolveModelProvider(event, ctx),
-    [SESSION_ATTRIBUTES.PI_MODEL_ID_CURRENT]: resolveModelId(event, ctx),
-    [SESSION_ATTRIBUTES.PI_THINKING_LEVEL_CURRENT]: resolveThinkingLevel(event, ctx),
+    [SESSION_ATTRIBUTES.PI_SESSION_VERSION]: readString(header, "version") ?? "unknown",
+    [SESSION_ATTRIBUTES.PI_MODEL_PROVIDER_CURRENT]: resolveModelProvider(ctx),
+    [SESSION_ATTRIBUTES.PI_MODEL_ID_CURRENT]: resolveModelId(ctx),
+    [SESSION_ATTRIBUTES.PI_THINKING_LEVEL_CURRENT]: resolveThinkingLevel(initialThinkingLevel),
     ...buildCommonSessionSpanAttributes(sessionId, config, lineage),
   });
 }
@@ -153,7 +171,7 @@ function createSessionStartHandler(
   loadConfigFn: LoadSessionConfig,
   startTelemetryFn: StartSessionTelemetry,
   options: RegisterHandlersOptions,
-): Handler {
+): PiHandler<"session_start"> {
   return handleSessionStart.bind(undefined, state, loadConfigFn, startTelemetryFn, options);
 }
 
@@ -162,19 +180,29 @@ async function handleSessionStart(
   loadConfigFn: LoadSessionConfig,
   startTelemetryFn: StartSessionTelemetry,
   options: RegisterHandlersOptions,
-  event: unknown,
-  ctx: ObservMeHandlerContext,
+  event: SessionStartEvent,
+  ctx: ExtensionContext,
 ): Promise<void> {
+  if (!(await resolvePendingTelemetryCleanupBeforeStart(state, ctx))) return;
+
   const previousSession = state.session;
-  if (previousSession) await shutDownPreviousSessionBeforeDuplicateStart(previousSession, ctx, state);
+  if (previousSession && !(await shutDownPreviousSessionBeforeDuplicateStart(previousSession, ctx, state))) return;
 
   await ensureProjectConfigForHandler(options, ctx);
   const loadedConfig = await loadSessionConfigForHandler(loadConfigFn, options, ctx);
   const config = loadedConfig.config;
-  await emitUnsafeCaptureWarning(config, ctx);
+  updateObsStatusRuntimeState({ config, configDiagnostics: loadedConfig.diagnostics });
+  clearObsStatusExportError();
 
+  if (!config.enabled) {
+    clearDisabledTelemetryRuntimeState(state);
+    await clearExtensionStatus(ctx);
+    return;
+  }
+
+  await emitUnsafeCaptureWarning(config, ctx);
   const recovery = await resolveStartupRecovery(event, ctx, config, options);
-  const recoveryCorrelation = recovery.customCorrelation ?? recovery.header?.correlation;
+  const recoveryCorrelation = recovery.customCorrelation;
   const lineage = createAgentLineageContext({
     config,
     env: buildRecoveryLineageEnv(config, recoveryCorrelation, options.env),
@@ -183,19 +211,31 @@ async function handleSessionStart(
       options.requireCompleteParentEnvelope ?? (options.trustedParentContext === true && recoveryCorrelation === undefined),
     failOpenInvalidPropagation: true,
   });
-  const session = await startTelemetryFn({
-    config,
-    lineage,
-    now: options.now,
-    wallClockNow: options.wallClockNow,
-  });
+  let session: ObservMeTelemetrySession;
+  try {
+    session = await startTelemetryFn({
+      config,
+      lineage,
+      now: options.now,
+      wallClockNow: options.wallClockNow,
+    });
+  } catch (error) {
+    await handleTelemetryStartupFailure(state, ctx, error);
+    return;
+  }
+
   session.now = options.now ?? session.now ?? monotonicNowMs;
   state.session = session;
 
   try {
-    updateObsStatusRuntimeState({ config: session.config, configDiagnostics: loadedConfig.diagnostics });
-    clearObsStatusExportError();
-    const attributes = buildSessionAttributes(event, ctx, session.config, lineage, recovery);
+    const attributes = buildSessionAttributes(
+      event,
+      ctx,
+      session.config,
+      lineage,
+      recovery,
+      options.getThinkingLevel?.(),
+    );
     const labels = metricLabels(session.config, lineage);
 
     session.sessionAttributes = attributes;
@@ -206,7 +246,7 @@ async function handleSessionStart(
     startObsSessionRuntimeState({
       sessionId: readString(attributes, SESSION_ATTRIBUTES.PI_SESSION_ID),
       traceId: readSpanTraceId(session.sessionSpan),
-      traceUrlTemplate: session.config.query.links.traceUrlTemplate,
+      config: session.config,
     });
     startObsAgentsRuntimeState({
       lineage,
@@ -218,29 +258,31 @@ async function handleSessionStart(
     session.metrics.sessionsStarted.add(1, labels);
     session.sessionSpan.addEvent(LOG_EVENT_NAMES.SESSION_STARTED, attributes);
     emitLifecycleLog(session.logger, LOG_EVENT_NAMES.SESSION_STARTED, attributes);
-    if (recovery.resumed && session.config.replayOnStart) emitStartupReplayTelemetry(session, attributes);
 
     if (isRootWorkflow(lineage)) {
       session.metrics.workflowsStarted.add(1, labels);
       emitLifecycleLog(session.logger, LOG_EVENT_NAMES.WORKFLOW_STARTED, attributes);
     }
 
-    await ctx.ui?.setStatus?.(EXTENSION_STATUS_KEY, EXTENSION_STATUS_VALUE);
+    ctx.ui?.setStatus?.(EXTENSION_STATUS_KEY, EXTENSION_STATUS_VALUE);
     activateSessionActiveAgent(session, labels);
+    if (session.config.agent.writeCorrelationEntry) {
+      appendSessionCorrelationEntry(options.appendEntry, lineage, recovery.customCorrelation);
+    }
   } catch (error) {
     await cleanUpFailedSessionStart(session, ctx, state);
     throw error;
   }
 }
 
-function createSessionShutdownHandler(state: HandlerSessionState): Handler {
+function createSessionShutdownHandler(state: HandlerSessionState): PiHandler<"session_shutdown"> {
   return handleSessionShutdown.bind(undefined, state);
 }
 
 async function handleSessionShutdown(
   state: HandlerSessionState,
-  event: unknown,
-  ctx: ObservMeHandlerContext,
+  event: SessionShutdownEvent,
+  ctx: ExtensionContext,
 ): Promise<void> {
   const session = state.session;
   if (!session) return;
@@ -256,8 +298,13 @@ async function resolveStartupRecovery(
 ): Promise<StartupRecoveryState> {
   const sessionFile = resolveSessionFilePath(event, ctx);
   const readHeader = options.readSessionHeader ?? readSessionHeaderFromFile;
-  const header = sessionFile ? await readHeader(sessionFile) : undefined;
-  const customCorrelation = config.agent.writeCorrelationEntry ? readExplicitCustomCorrelation(event) : undefined;
+  let header: SessionRecoveryHeader | undefined;
+  if (ctx.sessionManager) {
+    header = normalizeSessionHeader(ctx.sessionManager.getHeader());
+  } else if (sessionFile) {
+    header = await readHeader(sessionFile);
+  }
+  const customCorrelation = config.agent.writeCorrelationEntry ? readActiveBranchCorrelation(ctx) : undefined;
 
   return {
     resumed: isExistingSessionStart(event),
@@ -267,16 +314,6 @@ async function resolveStartupRecovery(
   };
 }
 
-function emitStartupReplayTelemetry(session: ObservMeTelemetrySession, attributes: AttributeMap): void {
-  const replayAttributes = {
-    ...attributes,
-    [COMMON_SPAN_ATTRIBUTES.OBSERVME_REPLAYED]: true,
-  };
-
-  session.sessionSpan?.addEvent(LOG_EVENT_NAMES.SESSION_STARTED, replayAttributes);
-  emitLifecycleLog(session.logger, LOG_EVENT_NAMES.SESSION_STARTED, replayAttributes);
-}
-
 function buildRecoveryLineageEnv(
   config: ObservMeConfig,
   correlation: MinimalSessionCorrelation | undefined,
@@ -291,19 +328,28 @@ function buildRecoveryLineageEnv(
     ...definedEnvValue(config.agent.parentIdEnv, correlation.parentAgentId),
     ...definedEnvValue(config.agent.rootIdEnv, correlation.rootAgentId),
     ...definedEnvValue(config.agent.parentSessionIdEnv, correlation.parentSessionId),
-    ...definedEnvValue(config.agent.depthEnv, correlation.depth === undefined ? undefined : String(correlation.depth)),
+    ...definedEnvValue(config.agent.depthEnv, recoveryPropagationDepth(correlation)),
     ...definedEnvValue(config.agent.spawnIdEnv, correlation.spawnId),
     ...definedEnvValue(config.agent.capabilityEnv, correlation.capability),
   };
 }
 
+function recoveryPropagationDepth(correlation: MinimalSessionCorrelation): string | undefined {
+  if (correlation.depth === undefined) return undefined;
+  const parentDepth = correlation.parentAgentId ? Math.max(0, correlation.depth - 1) : correlation.depth;
+  return String(parentDepth);
+}
+
 function definedEnvValue(name: string, value: string | undefined): NodeJS.ProcessEnv {
   return value === undefined || value === "" ? {} : { [name]: value };
 }
 
-function readExplicitCustomCorrelation(event: unknown): MinimalSessionCorrelation | undefined {
-  const value = readUnknown(event, "customCorrelation") ?? readUnknown(event, "observmeCorrelation");
-  return normalizeMinimalCorrelation(value);
+function readActiveBranchCorrelation(ctx: ObservMeHandlerContext): StartupRecoveryState["customCorrelation"] {
+  try {
+    return readLatestSessionCorrelation(ctx.sessionManager?.getBranch());
+  } catch {
+    return undefined;
+  }
 }
 
 function normalizeSessionHeader(value: unknown): SessionRecoveryHeader | undefined {
@@ -316,27 +362,9 @@ function normalizeSessionHeader(value: unknown): SessionRecoveryHeader | undefin
     timestamp: readString(value, "timestamp"),
     cwd: readString(value, "cwd"),
     parentSession: readString(value, "parentSession"),
-    correlation: normalizeMinimalCorrelation(readUnknown(value, "observmeCorrelation") ?? readUnknown(value, "correlation")),
   });
 }
 
-function normalizeMinimalCorrelation(value: unknown): MinimalSessionCorrelation | undefined {
-  if (!isRecord(value)) return undefined;
-
-  const correlation = withoutUndefinedObjectValues({
-    workflowId: readString(value, "workflowId"),
-    agentId: readString(value, "agentId"),
-    parentAgentId: readString(value, "parentAgentId"),
-    rootAgentId: readString(value, "rootAgentId"),
-    parentSessionId: readString(value, "parentSessionId"),
-    depth: readInteger(value, "depth"),
-    spawnId: readString(value, "spawnId"),
-    capability: readString(value, "capability"),
-  });
-
-  return Object.keys(correlation).length === 0 ? undefined : correlation;
-}
-
 function withoutUndefinedObjectValues<T extends Record<string, unknown>>(value: T): T {
   const definedEntries: Array<[string, unknown]> = [];
   for (const entry of Object.entries(value)) {
@@ -397,19 +425,84 @@ async function shutDownPreviousSessionBeforeDuplicateStart(
   session: ObservMeTelemetrySession,
   ctx: ObservMeHandlerContext,
   state: HandlerSessionState,
-): Promise<void> {
+): Promise<boolean> {
   emitLifecycleLog(session.logger, LOG_EVENT_NAMES.SESSION_DUPLICATE_START, buildDuplicateSessionStartAttributes(session));
 
   try {
-    await shutDownTelemetrySession(session, duplicateSessionStartShutdownEvent(), ctx, state);
+    const result = await shutDownTelemetrySession(session, duplicateSessionStartShutdownEvent(), ctx, state);
+    if (otelCleanupCompleted(result)) return true;
+    notifyPendingTelemetryCleanup(ctx);
+    return false;
   } catch (error) {
     recordDuplicateSessionStartShutdownError(session, error);
     clearObsSessionRuntimeState();
     clearObsAgentsRuntimeState();
     state.session = undefined;
+    state.pendingCleanup = session;
+    notifyPendingTelemetryCleanup(ctx);
+    return false;
   }
 }
 
+async function resolvePendingTelemetryCleanupBeforeStart(
+  state: HandlerSessionState,
+  ctx: ObservMeHandlerContext,
+): Promise<boolean> {
+  const pendingSession = state.pendingCleanup;
+  if (!pendingSession) return true;
+
+  const result = await recordControllerOperationResult(pendingSession, "shutdown");
+  if (otelCleanupCompleted(result)) {
+    if (state.pendingCleanup === pendingSession) state.pendingCleanup = undefined;
+    return true;
+  }
+
+  retainUnresolvedTelemetryCleanup(state, pendingSession, result);
+  notifyPendingTelemetryCleanup(ctx);
+  return false;
+}
+
+function retainUnresolvedTelemetryCleanup(
+  state: HandlerSessionState,
+  session: ObservMeTelemetrySession,
+  result: BoundedOtelOperationResult | undefined,
+): void {
+  if (!result || otelCleanupCompleted(result)) return;
+
+  state.pendingCleanup = session;
+  void result.settlement?.then(releaseSettledTelemetryCleanup.bind(undefined, state, session));
+}
+
+function releaseSettledTelemetryCleanup(
+  state: HandlerSessionState,
+  session: ObservMeTelemetrySession,
+  settlement: OtelOperationSettlement,
+): void {
+  if (!otelCleanupCompleted(settlement) || state.pendingCleanup !== session) return;
+  state.pendingCleanup = undefined;
+}
+
+function otelCleanupCompleted(
+  result: Pick<BoundedOtelOperationResult, "completed" | "timedOut" | "error">,
+): boolean {
+  return result.completed && !result.timedOut && !result.error;
+}
+
+function notifyPendingTelemetryCleanup(ctx: ObservMeHandlerContext): void {
+  if (ctx.hasUI === false || !ctx.ui?.notify) return;
+
+  const message = "ObservMe telemetry startup was deferred because prior OTEL shutdown cleanup is still unresolved.";
+  try {
+    void Promise.resolve(ctx.ui.notify(message, "warning")).catch(ignorePendingCleanupDiagnosticError);
+  } catch {
+    return;
+  }
+}
+
+function ignorePendingCleanupDiagnosticError(): undefined {
+  return undefined;
+}
+
 function buildDuplicateSessionStartAttributes(session: ObservMeTelemetrySession): AttributeMap {
   return withoutUndefinedAttributes({
     [LOG_ATTRIBUTES.PI_SESSION_ID]: readString(session.sessionAttributes, SESSION_ATTRIBUTES.PI_SESSION_ID),
@@ -444,6 +537,7 @@ async function cleanUpFailedSessionStart(
   state: HandlerSessionState,
 ): Promise<void> {
   const labels = metricLabels(session.config, session.lineage);
+  let shutdownResult: BoundedOtelOperationResult | undefined;
 
   try {
     deactivateSessionActiveAgent(session, labels);
@@ -451,10 +545,11 @@ async function cleanUpFailedSessionStart(
     endActiveSpan(session, session.sessionSpan);
     await clearExtensionStatus(ctx);
     await recordControllerOperationResult(session, "flush");
-    await recordControllerOperationResult(session, "shutdown");
+    shutdownResult = await recordControllerOperationResult(session, "shutdown");
   } finally {
     disposeSessionActiveAgentLease(session);
     clearTelemetrySessionRuntimeState(session, state);
+    retainUnresolvedTelemetryCleanup(state, session, shutdownResult);
   }
 }
 
@@ -463,10 +558,11 @@ async function shutDownTelemetrySession(
   event: unknown,
   ctx: ObservMeHandlerContext,
   state: HandlerSessionState,
-): Promise<void> {
+): Promise<BoundedOtelOperationResult> {
   const labels = metricLabels(session.config, session.lineage);
   const shutdownAttributes = buildShutdownAttributes(event, session);
   const failed = workflowFailed(event);
+  let shutdownResult: BoundedOtelOperationResult | undefined;
 
   try {
     deactivateSessionActiveAgent(session, labels);
@@ -478,10 +574,12 @@ async function shutDownTelemetrySession(
     endActiveSpan(session, session.sessionSpan);
     await clearExtensionStatus(ctx);
     await recordControllerOperationResult(session, "flush");
-    await recordControllerOperationResult(session, "shutdown");
+    shutdownResult = await recordControllerOperationResult(session, "shutdown");
+    return shutdownResult;
   } finally {
     disposeSessionActiveAgentLease(session);
     clearTelemetrySessionRuntimeState(session, state);
+    retainUnresolvedTelemetryCleanup(state, session, shutdownResult);
   }
 }
 
@@ -520,6 +618,39 @@ function clearTelemetrySessionRuntimeState(
   if (state.session === session) state.session = undefined;
 }
 
+function clearDisabledTelemetryRuntimeState(state: HandlerSessionState): void {
+  state.session = undefined;
+  clearObsSessionRuntimeState();
+  clearObsAgentsRuntimeState();
+}
+
+async function handleTelemetryStartupFailure(
+  state: HandlerSessionState,
+  ctx: ObservMeHandlerContext,
+  error: unknown,
+): Promise<void> {
+  clearDisabledTelemetryRuntimeState(state);
+  if (!(error instanceof ObservMeOtelStartupError)) throw error;
+
+  recordObsStatusExportResult({ operation: "startup", error });
+  await clearExtensionStatus(ctx);
+  notifyTelemetryStartupFailure(ctx, error.message);
+}
+
+function notifyTelemetryStartupFailure(ctx: ObservMeHandlerContext, message: string): void {
+  if (ctx.hasUI === false || !ctx.ui?.notify) return;
+
+  try {
+    void Promise.resolve(ctx.ui.notify(message, "error")).catch(ignoreStartupDiagnosticError);
+  } catch {
+    return;
+  }
+}
+
+function ignoreStartupDiagnosticError(): undefined {
+  return undefined;
+}
+
 async function clearExtensionStatus(ctx: ObservMeHandlerContext): Promise<void> {
   try {
     await ctx.ui?.setStatus?.(EXTENSION_STATUS_KEY, undefined);
@@ -531,10 +662,11 @@ async function clearExtensionStatus(ctx: ObservMeHandlerContext): Promise<void>
 async function recordControllerOperationResult(
   session: ObservMeTelemetrySession,
   operation: BoundedOtelOperationResult["operation"],
-): Promise<void> {
+): Promise<BoundedOtelOperationResult> {
   const result = await runControllerOperation(session, operation);
   recordObsStatusExportResult(result);
   recordExportOperationResult(session, result);
+  return result;
 }
 
 async function runControllerOperation(

package/src/pi/event-handlers/llm.ts

@@ -1,4 +1,8 @@
 import { SpanStatusCode } from "@opentelemetry/api";
+import type {
+  BeforeProviderRequestEvent,
+  ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
 import { recordObsSessionLlmCall } from "../../commands/obs-session.ts";
 import { LOG_EVENT_NAMES } from "../../semconv/metrics.ts";
 import { SPAN_NAMES } from "../../semconv/spans.ts";
@@ -26,7 +30,7 @@ import {
   startActiveChildSpan,
 } from "../handler-internals.ts";
 import type { HandlerRegistrar } from "../handler-runtime.ts";
-import type { Handler, HandlerSessionState, ObservMeHandlerContext } from "../handler-types.ts";
+import type { HandlerSessionState, PiEvent, PiHandler } from "../handler-types.ts";
 import { recordBashExecution } from "./tool-bash.ts";
 
 export function registerLlmHandlers(registrar: HandlerRegistrar, state: HandlerSessionState): void {
@@ -35,14 +39,14 @@ export function registerLlmHandlers(registrar: HandlerRegistrar, state: HandlerS
   registrar.add("message_end", createMessageEndHandler(state));
 }
 
-function createBeforeProviderRequestHandler(state: HandlerSessionState): Handler {
+function createBeforeProviderRequestHandler(state: HandlerSessionState): PiHandler<"before_provider_request"> {
   return handleBeforeProviderRequest.bind(undefined, state);
 }
 
 function handleBeforeProviderRequest(
   state: HandlerSessionState,
-  event: unknown,
-  ctx: ObservMeHandlerContext,
+  event: BeforeProviderRequestEvent,
+  ctx: ExtensionContext,
 ): void {
   const session = state.session;
   if (!session) return;
@@ -61,14 +65,14 @@ function handleBeforeProviderRequest(
   emitLifecycleLog(session.logger, LOG_EVENT_NAMES.LLM_REQUEST_STARTED, attributes);
 }
 
-function createAfterProviderResponseHandler(state: HandlerSessionState): Handler {
+function createAfterProviderResponseHandler(state: HandlerSessionState): PiHandler<"after_provider_response"> {
   return handleAfterProviderResponse.bind(undefined, state);
 }
 
 function handleAfterProviderResponse(
   state: HandlerSessionState,
-  event: unknown,
-  _ctx: ObservMeHandlerContext,
+  event: PiEvent<"after_provider_response">,
+  _ctx: ExtensionContext,
 ): void {
   const session = state.session;
   if (!session) return;
@@ -77,11 +81,15 @@ function handleAfterProviderResponse(
   span?.setAttributes(buildLlmResponseAttributes(event));
 }
 
-function createMessageEndHandler(state: HandlerSessionState): Handler {
+function createMessageEndHandler(state: HandlerSessionState): PiHandler<"message_end"> {
   return handleMessageEnd.bind(undefined, state);
 }
 
-function handleMessageEnd(state: HandlerSessionState, event: unknown, _ctx: ObservMeHandlerContext): void {
+function handleMessageEnd(
+  state: HandlerSessionState,
+  event: PiEvent<"message_end">,
+  _ctx: ExtensionContext,
+): void {
   const session = state.session;
   if (!session) return;