nemo-relay-openclaw 0.3.0 → 0.5.0

package/src/hook-replay/session.ts

@@ -9,7 +9,10 @@
  * those identifiers and owns the root `openclaw.session` scope lifecycle.
  */
 import type { NemoRelayHookBackendConfig } from '../config.js';
-import { evictExpiredRecords, tupleKey as tupleKeyFromCorrelation } from './correlation.js';
+import {
+  evictExpiredRecords,
+  tupleKey as tupleKeyFromCorrelation,
+} from './correlation.js';
 import type {
   PluginHookAgentContext,
   PluginHookLlmOutputEvent,
@@ -18,6 +21,7 @@ import type {
 import type { PluginLogger } from 'openclaw/plugin-sdk/plugin-entry';
 import type { JsonObject as JsonRecord } from 'nemo-relay-node/typed';
 import type { NemoRelayRuntimeModule } from '../modules.js';
+import { toJsonRecord } from './marks.js';
 
 export type SessionLookupInput = {
   sessionId?: string | undefined;
@@ -29,12 +33,17 @@ export type SessionLookupInput = {
 
 export type EnsureSessionInput = SessionLookupInput & {
   agentId?: string | undefined;
+  parentHandle?: ReturnType<NemoRelayRuntimeModule['pushScope']> | undefined;
+  scopeRole?: 'subagent' | undefined;
   source: 'session_start' | 'lazy_session';
   resumedFrom?: string | undefined;
   timestamp?: number | undefined;
+  deferRootOpen?: boolean | undefined;
 };
 
 export type SessionState = {
+  /** Immutable internal owner key used for replay buffers and alias lookups. */
+  ownerKey: string;
   sessionId: string;
   sessionKey?: string;
   agentId?: string;
@@ -51,10 +60,15 @@ export type SessionState = {
   assistantMessageWrites?: AssistantMessageRecord[];
   stack: ReturnType<NemoRelayRuntimeModule['createScopeStack']>;
   rootHandle?: ReturnType<NemoRelayRuntimeModule['pushScope']>;
+  scopeRole?: 'subagent';
+  pendingRootOpen?: boolean;
+  pendingRootTimestampMicros?: number;
+  pendingRootRunId?: string;
+  pendingCapturedEmits?: Array<{ label: string; emit: () => void }>;
 };
 
 export type PendingLlmOutputRecord = {
-  sessionKey: string;
+  sessionOwnerKey: string;
   sessionId: string;
   runId: string;
   provider: string;
@@ -66,7 +80,7 @@ export type PendingLlmOutputRecord = {
 };
 
 export type LlmInputRecord = {
-  sessionKey: string;
+  sessionOwnerKey: string;
   sessionId: string;
   runId: string;
   provider: string;
@@ -80,7 +94,7 @@ export type LlmInputRecord = {
 };
 
 export type AssistantMessageRecord = {
-  sessionKey: string;
+  sessionOwnerKey: string;
   provider: string;
   model: string;
   assistantTexts: string[];
@@ -93,7 +107,7 @@ export type AssistantMessageRecord = {
 };
 
 export type ModelCallRecord = {
-  sessionKey: string;
+  sessionOwnerKey: string;
   runId: string;
   callId: string;
   provider: string;
@@ -144,8 +158,32 @@ export type SessionManager = {
   replayPendingLlmOutputsForSession: (session: SessionState, options: { allowPlaceholderRequest: boolean }) => void;
   emitUnpairedModelCallTimingMarks: (session: SessionState) => void;
   logBoundedWarn: (key: string, message: string) => void;
+  resolveSessionRootContext?: (input: EnsureSessionInput) => Partial<EnsureSessionInput> | undefined;
 };
 
+/** Merge lineage/root-context defaults without letting undefined hook fields erase them. */
+function mergeEnsureSessionInput(
+  input: EnsureSessionInput,
+  context: Partial<EnsureSessionInput> | undefined,
+): EnsureSessionInput {
+  if (!context) {
+    return input;
+  }
+
+  const merged: Record<string, unknown> = { ...context, ...input };
+  for (const [key, value] of Object.entries(context)) {
+    if (merged[key] === undefined && value !== undefined) {
+      merged[key] = value;
+    }
+  }
+  return merged as EnsureSessionInput;
+}
+
+/** Return the session key that belongs to the session itself, not a requester alias. */
+function ownSessionKey(input: SessionLookupInput): string | undefined {
+  return input.sessionKey ?? input.childSessionKey;
+}
+
 /** Return all keys that may identify an existing OpenClaw session. */
 export function lookupSessionKeys(input: SessionLookupInput): string[] {
   return [input.sessionId, input.sessionKey, input.requesterSessionKey, input.childSessionKey, input.runId].filter(
@@ -153,15 +191,15 @@ export function lookupSessionKeys(input: SessionLookupInput): string[] {
   );
 }
 
-/** Return keys that should alias to a canonical session once it is known. */
+/** Return keys that should alias to one stable internal session owner. */
 export function aliasSessionKeys(input: SessionLookupInput): string[] {
-  return [input.sessionId, input.sessionKey, input.requesterSessionKey, input.runId].filter(
+  return [input.sessionId, input.sessionKey, input.childSessionKey, input.requesterSessionKey, input.runId].filter(
     (value): value is string => typeof value === 'string' && value.length > 0,
   );
 }
 
-/** Resolve a hook's session identity to the canonical session id used in replay state. */
-export function resolveSessionKey(state: HookReplayBackendState, input: SessionLookupInput): string | undefined {
+/** Resolve a hook's session identity to the stable owner key used in replay state. */
+export function resolveSessionOwnerKey(state: HookReplayBackendState, input: SessionLookupInput): string | undefined {
   for (const key of lookupSessionKeys(input)) {
     const canonical = state.sessionAliases.get(key);
     if (canonical) {
@@ -172,14 +210,14 @@ export function resolveSessionKey(state: HookReplayBackendState, input: SessionL
   return input.sessionId ?? input.sessionKey ?? input.childSessionKey ?? input.runId;
 }
 
-/** Remember equivalent hook identifiers so later events attach to the same root span. */
+/** Remember equivalent hook identifiers so later events attach to the same owner and root span. */
 export function rememberSessionAliases(
   state: HookReplayBackendState,
   session: SessionState,
   input: SessionLookupInput,
 ): void {
   for (const alias of aliasSessionKeys(input)) {
-    state.sessionAliases.set(alias, session.sessionId);
+    state.sessionAliases.set(alias, session.ownerKey);
   }
 }
 
@@ -204,47 +242,102 @@ export function createHookReplayState(): HookReplayBackendState {
 
 /** Return an existing session or lazily create a root session scope for replay. */
 export function ensureSession(manager: SessionManager, input: EnsureSessionInput): SessionState | undefined {
-  const key = resolveSessionKey(manager.state, input);
-  if (!key) {
+  const resolvedInput = mergeEnsureSessionInput(input, manager.resolveSessionRootContext?.(input));
+  const ownerKey = resolveSessionOwnerKey(manager.state, resolvedInput);
+  if (!ownerKey) {
     manager.state.counters.skippedEvents += 1;
     manager.logBoundedWarn('missing-session-key', 'nemo-relay skipped replay because no session/run key was available');
     return undefined;
   }
 
-  const existing = manager.state.sessions.get(key);
+  const existing = manager.state.sessions.get(ownerKey);
   if (existing) {
-    rememberSessionAliases(manager.state, existing, input);
+    enrichSession(existing, resolvedInput);
+    rememberSessionAliases(manager.state, existing, resolvedInput);
+    if (!existing.rootHandle && resolvedInput.deferRootOpen !== true) {
+      materializeSessionRoot(manager, existing, resolvedInput);
+    }
     return existing;
   }
 
-  const canonicalSessionId = input.sessionId ?? key;
-  const aliased = manager.state.sessions.get(canonicalSessionId);
-  if (aliased) {
-    rememberSessionAliases(manager.state, aliased, input);
-    return aliased;
-  }
-
   const stack = manager.nf.createScopeStack();
   const session: SessionState = {
-    sessionId: canonicalSessionId,
-    source: input.source,
+    ownerKey,
+    sessionId: resolvedInput.sessionId ?? ownerKey,
+    source: resolvedInput.source,
     stack,
   };
 
-  if (input.sessionKey !== undefined) {
-    session.sessionKey = input.sessionKey;
+  const sessionKey = ownSessionKey(resolvedInput);
+  if (sessionKey !== undefined) {
+    session.sessionKey = sessionKey;
+  }
+  if (resolvedInput.agentId !== undefined) {
+    session.agentId = resolvedInput.agentId;
   }
-  if (input.agentId !== undefined) {
+  if (resolvedInput.resumedFrom !== undefined) {
+    session.resumedFrom = resolvedInput.resumedFrom;
+  }
+  if (resolvedInput.scopeRole !== undefined) {
+    session.scopeRole = resolvedInput.scopeRole;
+  }
+
+  if (resolvedInput.deferRootOpen === true) {
+    session.pendingRootOpen = true;
+    if (resolvedInput.timestamp !== undefined) {
+      session.pendingRootTimestampMicros = resolvedInput.timestamp;
+    }
+    if (resolvedInput.runId !== undefined) {
+      session.pendingRootRunId = resolvedInput.runId;
+    }
+    session.pendingCapturedEmits = [];
+  } else {
+    materializeSessionRoot(manager, session, resolvedInput);
+  }
+  manager.state.sessions.set(session.ownerKey, session);
+  rememberSessionAliases(manager.state, session, resolvedInput);
+  return session;
+}
+
+/** Fill stable session identifiers from later hooks without clobbering established values. */
+function enrichSession(session: SessionState, input: EnsureSessionInput): void {
+  const sessionKey = ownSessionKey(input);
+  if (session.sessionKey === undefined && sessionKey !== undefined) {
+    session.sessionKey = sessionKey;
+  }
+  if (session.agentId === undefined && input.agentId !== undefined) {
     session.agentId = input.agentId;
   }
-  if (input.resumedFrom !== undefined) {
+  if (session.resumedFrom === undefined && input.resumedFrom !== undefined) {
     session.resumedFrom = input.resumedFrom;
   }
+  if (session.scopeRole === undefined && input.scopeRole !== undefined) {
+    session.scopeRole = input.scopeRole;
+  }
+  if (!session.rootHandle && input.sessionId !== undefined) {
+    session.sessionId = input.sessionId;
+  }
+  if (!session.rootHandle && session.source === 'lazy_session' && input.source === 'session_start') {
+    session.source = 'session_start';
+  }
+  if (!session.rootHandle && input.timestamp !== undefined) {
+    if (input.source === 'session_start' || session.pendingRootTimestampMicros === undefined) {
+      session.pendingRootTimestampMicros = input.timestamp;
+    }
+  }
+  if (!session.rootHandle && session.pendingRootRunId === undefined && input.runId !== undefined) {
+    session.pendingRootRunId = input.runId;
+  }
+}
 
-  openSessionRoot(manager, session, input);
-  manager.state.sessions.set(session.sessionId, session);
-  rememberSessionAliases(manager.state, session, input);
-  return session;
+/** Queue an emit callback until a deferred session root can open with honest lineage. */
+export function queueCapturedEmit(session: SessionState, label: string, emit: () => void): boolean {
+  if (session.rootHandle || session.pendingRootOpen !== true) {
+    return false;
+  }
+
+  (session.pendingCapturedEmits ??= []).push({ label, emit });
+  return true;
 }
 
 /** Flush pending LLM output/timing state before the root session closes. */
@@ -261,6 +354,7 @@ export function closeSessionRoot(
   session: SessionState,
   summary: JsonRecord,
   rootOutput: JsonRecord = summary,
+  metadata?: JsonRecord | null,
   timestamp?: number,
 ): void {
   manager.emitCapturedUnderSession('session_end', session, () => {
@@ -268,7 +362,7 @@ export function closeSessionRoot(
       return;
     }
 
-    manager.nf.event('openclaw.session_end', session.rootHandle, summary, null, timestamp ?? null);
+    manager.nf.event('openclaw.session_end', session.rootHandle, summary, metadata ?? null, timestamp ?? null);
     manager.state.counters.marksEmitted += 1;
     manager.nf.popScope(session.rootHandle, rootOutput, timestamp ?? null);
     delete session.rootHandle;
@@ -277,7 +371,7 @@ export function closeSessionRoot(
 
 /** Remove a closed session from active replay state. */
 export function deleteSession(state: HookReplayBackendState, session: SessionState): void {
-  state.sessions.delete(session.sessionId);
+  state.sessions.delete(session.ownerKey);
 }
 
 /** Insert a correlation record while bounding retained entries per key. */
@@ -303,38 +397,69 @@ export function evictExpiredCorrelationRecords(state: HookReplayBackendState, no
   evictExpiredRecords(state.modelTimingsByLlmKey, nowMs, ttlMs);
 }
 
-/** Open the root NeMo Relay scope for one OpenClaw session and emit session_start. */
-function openSessionRoot(manager: SessionManager, session: SessionState, input: EnsureSessionInput): void {
+/** Open a deferred or new root session scope and flush queued child emissions. */
+export function materializeSessionRoot(manager: SessionManager, session: SessionState, input: EnsureSessionInput): void {
+  if (session.rootHandle) {
+    return;
+  }
+
+  enrichSession(session, input);
+
+  const timestampMicros = input.timestamp ?? session.pendingRootTimestampMicros ?? null;
+  const rootRunId = input.runId ?? session.pendingRootRunId;
+
   const data: JsonRecord = {
     sessionId: session.sessionId,
     source: session.source,
     ...(session.sessionKey === undefined ? {} : { sessionKey: session.sessionKey }),
     ...(session.agentId === undefined ? {} : { agentId: session.agentId }),
-    ...(input.runId === undefined ? {} : { runId: input.runId }),
+    ...(rootRunId === undefined ? {} : { runId: rootRunId }),
     ...(session.resumedFrom === undefined ? {} : { resumedFrom: session.resumedFrom }),
   };
+  const metadata = toJsonRecord({
+    source: session.source === 'session_start' ? 'openclaw.session_start' : 'openclaw.lazy_session',
+    hook_event_name: session.source === 'session_start' ? 'session_start' : undefined,
+    sessionId: session.sessionId,
+    sessionKey: session.sessionKey,
+    agentId: session.agentId,
+    runId: rootRunId,
+    nemo_relay_scope_role: session.scopeRole,
+  });
 
-  manager.emitCapturedUnderSession('session_start', session, () => {
+  const previousStack = manager.nf.currentScopeStack();
+  try {
+    manager.nf.setThreadScopeStack(session.stack);
     session.rootHandle = manager.nf.pushScope(
       'openclaw.session',
       agentScopeType(manager.nf),
-      null,
+      input.parentHandle ?? null,
       null,
       data,
-      null,
-      null,
-      input.timestamp ?? null,
+      metadata,
+      data,
+      timestampMicros,
     );
-    manager.nf.event('openclaw.session_start', session.rootHandle, data, null, input.timestamp ?? null);
+    manager.nf.event('openclaw.session_start', session.rootHandle, data, metadata, timestampMicros);
     manager.state.counters.marksEmitted += 1;
-  });
+  } finally {
+    manager.nf.setThreadScopeStack(previousStack);
+  }
+
+  delete session.pendingRootOpen;
+  delete session.pendingRootTimestampMicros;
+  delete session.pendingRootRunId;
+  const pendingEmits = session.pendingCapturedEmits ?? [];
+  delete session.pendingCapturedEmits;
+  for (const pending of pendingEmits) {
+    manager.emitCapturedUnderSession(pending.label, session, pending.emit);
+  }
 }
 
 /** Cancel timers that would otherwise replay late LLM outputs after session close. */
 function cancelPendingLlmOutputTimers(state: HookReplayBackendState, session: SessionState): void {
   for (const records of state.llmOutputsPendingInput.values()) {
     for (const record of records) {
-      if (record.sessionKey === session.sessionId && record.timer) {
+      if (record.sessionOwnerKey === session.ownerKey && record.timer) {
         clearTimeout(record.timer);
         record.timer = undefined;
       }
@@ -344,22 +469,22 @@ function cancelPendingLlmOutputTimers(state: HookReplayBackendState, session: Se
 
 /** Remove all correlation records and aliases owned by a closed session. */
 function evictSessionCorrelationRecords(state: HookReplayBackendState, session: SessionState): void {
-  evictFromRecordMap(state.llmInputs, session.sessionId);
-  evictFromRecordMap(state.llmOutputsPendingInput, session.sessionId);
-  evictFromRecordMap(state.modelCallsByCallId, session.sessionId);
-  evictFromRecordMap(state.modelTimingsByLlmKey, session.sessionId);
+  evictFromRecordMap(state.llmInputs, session.ownerKey);
+  evictFromRecordMap(state.llmOutputsPendingInput, session.ownerKey);
+  evictFromRecordMap(state.modelCallsByCallId, session.ownerKey);
+  evictFromRecordMap(state.modelTimingsByLlmKey, session.ownerKey);
 
-  for (const [alias, canonical] of state.sessionAliases) {
-    if (canonical === session.sessionId || alias === session.sessionId) {
+  for (const [alias, ownerKey] of state.sessionAliases) {
+    if (ownerKey === session.ownerKey) {
       state.sessionAliases.delete(alias);
     }
   }
 }
 
 /** Drop records for one session from a single keyed correlation map. */
-function evictFromRecordMap<T extends { sessionKey: string }>(map: Map<string, T[]>, sessionKey: string): void {
+function evictFromRecordMap<T extends { sessionOwnerKey: string }>(map: Map<string, T[]>, ownerKey: string): void {
   for (const [key, records] of map) {
-    const retained = records.filter((record) => record.sessionKey !== sessionKey);
+    const retained = records.filter((record) => record.sessionOwnerKey !== ownerKey);
     if (retained.length === 0) {
       map.delete(key);
     } else {

package/src/hook-replay/tool.ts

@@ -22,12 +22,14 @@ export async function guardBeforeToolCall(
   event: PluginHookBeforeToolCallEvent,
   ctx: PluginHookToolContext,
 ): Promise<void> {
+  const observedAtMicros = nowMicros();
   const session = ensureSession(manager, {
     sessionId: ctx.sessionId,
     sessionKey: ctx.sessionKey,
     runId: event.runId ?? ctx.runId,
     agentId: ctx.agentId,
     source: 'lazy_session',
+    timestamp: observedAtMicros,
   });
   const args = toJsonValue(event.params ?? {});
 
@@ -50,12 +52,15 @@ export function replayAfterToolCall(
   event: PluginHookAfterToolCallEvent,
   ctx: PluginHookToolContext,
 ): void {
+  const endMicros = nowMicros();
+  const sessionTimestamp = startMicrosFromDuration(endMicros, event.durationMs) ?? endMicros;
   const session = ensureSession(manager, {
     sessionId: ctx.sessionId,
     sessionKey: ctx.sessionKey,
     runId: event.runId ?? ctx.runId,
     agentId: ctx.agentId,
     source: 'lazy_session',
+    timestamp: sessionTimestamp,
   });
 
   const blockedDetails = blockedToolDetails(event, { runId: event.runId ?? ctx.runId });
@@ -67,6 +72,16 @@ export function replayAfterToolCall(
         session,
         name: 'openclaw.tool_blocked',
         data: blockedDetails,
+        metadata: toJsonRecord({
+          source: 'openclaw.after_tool_call',
+          hook_event_name: 'after_tool_call',
+          sessionId: session.sessionId,
+          sessionKey: session.sessionKey,
+          agentId: session.agentId,
+          runId: event.runId ?? ctx.runId,
+          toolCallId: event.toolCallId ?? ctx.toolCallId,
+        }),
+        timestamp: endMicros,
       });
     });
     return;
@@ -76,12 +91,12 @@ export function replayAfterToolCall(
     return;
   }
 
-  const endMicros = nowMicros();
   const metadata = toJsonRecord({
     source: 'openclaw.after_tool_call',
     runId: event.runId ?? ctx.runId,
-    sessionId: ctx.sessionId,
-    sessionKey: ctx.sessionKey,
+    sessionId: session.sessionId,
+    sessionKey: session.sessionKey,
+    agentId: session.agentId,
     toolCallId: event.toolCallId ?? ctx.toolCallId,
     durationMs: event.durationMs,
   });