diff --git a/src/main/agent-runtime/inbox-projection.ts b/src/main/agent-runtime/inbox-projection.ts index 1f426fc1..904dd79f 100644 --- a/src/main/agent-runtime/inbox-projection.ts +++ b/src/main/agent-runtime/inbox-projection.ts @@ -66,12 +66,20 @@ function resolveSenderLabel( /** * Projects a stored `AgentMessageRecord` into a pi `AgentMessage` * (role: "user") suitable for appending to the recipient's session - * transcript. `timestamp` carries the message's `createdAt` in - * ms-since-epoch so transcript ordering aligns with wall-clock time. + * transcript. + * + * `timestampOverride`, when supplied, replaces the wall-clock timestamp + * derived from `record.createdAt`. Callers use this to allocate a + * monotonically unique value per run — two inbound messages landing in + * the same millisecond would otherwise share the same key, and the + * `RuntimeBridge` correlation index would lose one author's identity + * on the `message_end` callback. The actual value still lives in pi's + * `timestamp` number so ordering semantics remain identical. */ export function projectInboxMessageToPi( record: AgentMessageRecord, directory: AgentDirectory, + options?: { timestampOverride?: number }, ): AgentMessage { const senderLabel = resolveSenderLabel(record, directory); const channelLabel = formatChannelLabel(record); @@ -87,10 +95,14 @@ export function projectInboxMessageToPi( if (body.length > 0) parts.push(body); const text = parts.join("\n\n"); - const timestamp = Date.parse(record.createdAt); + let timestamp = options?.timestampOverride; + if (typeof timestamp !== "number" || !Number.isFinite(timestamp)) { + const parsed = Date.parse(record.createdAt); + timestamp = Number.isFinite(parsed) ? parsed : Date.now(); + } return { role: "user", content: [{ type: "text", text }], - timestamp: Number.isFinite(timestamp) ? timestamp : Date.now(), + timestamp, } as AgentMessage; } diff --git a/src/main/agent-runtime/runtime-bridge.ts b/src/main/agent-runtime/runtime-bridge.ts index 10fe1322..9b7a9940 100644 --- a/src/main/agent-runtime/runtime-bridge.ts +++ b/src/main/agent-runtime/runtime-bridge.ts @@ -60,6 +60,19 @@ interface InboundCorrelationEntry { fromAgentId: string; senderKind: AgentMessageRecord["senderKind"]; messageId: string; + /** + * Stable pi-message timestamp that the projected user turn carries. + * Reverse-indexed so `message_end` (which only surfaces the pi + * message shape) can resolve back to `messageId` without relying on + * wall-clock millisecond precision. + */ + piTimestamp: number; +} + +interface InboundCorrelationRegistry { + byMessageId: Map; + byPiTimestamp: Map; + lastAllocatedTimestamp: number; } export class RuntimeBridge { @@ -67,15 +80,16 @@ export class RuntimeBridge { private readonly activeRuns = new Map(); private readonly sessionStores = new Map(); /** - * Per-run mapping from pi-message timestamp to the inbound message that - * sourced it. The transcript only carries a rendered `[Message from X]` - * header, which is insufficient to recover the original author when the - * thread-event writer fires on `message_end`. Writers (inbox projection - * during prepareRun, steer injection during the run) register the entry - * before pi surfaces the message so the author-id survives the - * round-trip. + * Per-run inbound-message registry. Primary identity is + * `record.id` (the `agent_messages.id`). A monotonically allocated + * `piTimestamp` lives on every projected user turn so + * `message_end` — which only surfaces the pi-message shape — can + * invert the mapping back to the real author. The wall-clock + * `createdAt` is not enough: two messages that land in the same + * millisecond would collide on a timestamp-keyed map and one peer's + * identity would be silently lost in thread-event projection. */ - private readonly inboundCorrelations = new Map>(); + private readonly inboundCorrelations = new Map(); constructor(private readonly deps: RuntimeBridgeDeps) { this.costWriter = new CostWriter({ db: deps.db }); @@ -131,25 +145,45 @@ export class RuntimeBridge { } /** - * Register an inbound message that is about to enter (or has already - * entered) the recipient's transcript as a user-role turn. Called from - * the prepareRun inbox projection and from `RuntimeMessageInjector` on - * steer-ack success so the `recordThreadEvent` hook can recover the - * original author instead of labelling peer messages as the recipient. + * Build the pi projection for an inbound message and register the + * correlation under `record.id`. Returns the projected user turn so + * callers (the prepareRun inbox loop, the runtime message injector) + * append / steer the same object that the bridge has indexed. + * + * A monotonically increasing `piTimestamp` is stamped onto the + * returned message so two inbound rows that share the same + * wall-clock millisecond still resolve to distinct correlation + * entries when `message_end` fires. */ - registerInboundInjection(runId: string, record: AgentMessageRecord): void { - const timestamp = Date.parse(record.createdAt); - if (!Number.isFinite(timestamp)) return; - let runMap = this.inboundCorrelations.get(runId); - if (!runMap) { - runMap = new Map(); - this.inboundCorrelations.set(runId, runMap); - } - runMap.set(timestamp, { + projectInboundMessage(runId: string, record: AgentMessageRecord): AgentMessage { + const registry = this.getOrCreateInboundRegistry(runId); + const base = Date.parse(record.createdAt); + const baseline = Number.isFinite(base) ? base : Date.now(); + const allocated = Math.max(baseline, registry.lastAllocatedTimestamp + 1); + registry.lastAllocatedTimestamp = allocated; + const entry: InboundCorrelationEntry = { fromAgentId: record.fromAgentId, senderKind: record.senderKind, messageId: record.id, - }); + piTimestamp: allocated, + }; + registry.byMessageId.set(record.id, entry); + registry.byPiTimestamp.set(allocated, record.id); + return projectInboxMessageToPi(record, { + getAgentName: (id) => this.resolveAgentName(id), + }, { timestampOverride: allocated }); + } + + private getOrCreateInboundRegistry(runId: string): InboundCorrelationRegistry { + const existing = this.inboundCorrelations.get(runId); + if (existing) return existing; + const registry: InboundCorrelationRegistry = { + byMessageId: new Map(), + byPiTimestamp: new Map(), + lastAllocatedTimestamp: 0, + }; + this.inboundCorrelations.set(runId, registry); + return registry; } /** Returns the active run context for a given agent, if any. Used by @@ -373,9 +407,11 @@ export class RuntimeBridge { timestamp: number | undefined, ): InboundCorrelationEntry | null { if (typeof timestamp !== "number" || !Number.isFinite(timestamp)) return null; - const runMap = this.inboundCorrelations.get(runId); - if (!runMap) return null; - return runMap.get(timestamp) ?? null; + const registry = this.inboundCorrelations.get(runId); + if (!registry) return null; + const messageId = registry.byPiTimestamp.get(timestamp); + if (!messageId) return null; + return registry.byMessageId.get(messageId) ?? null; } /** @@ -495,15 +531,16 @@ export class RuntimeBridge { // whole closure. const store = this.getOrCreateStore(payload.taskSessionId); const inboxRecords = this.deps.db.listAgentInboxUnread(payload.companyId, payload.agentId); - const acceptedRecords: AgentMessageRecord[] = []; if (inboxRecords.length > 0) { this.deps.db.withTransaction(() => { for (const record of inboxRecords) { const claimed = this.deps.db.inboxProjections.claimProjection(record.id, payload.taskSessionId); if (!claimed) continue; - const projected = projectInboxMessageToPi(record, { - getAgentName: (id) => this.resolveAgentName(id), - }); + // `projectInboundMessage` allocates a monotonic pi-timestamp + // and indexes the correlation under `record.id`. Two rows + // that share the same wall-clock millisecond still resolve + // to distinct authors at `message_end`. + const projected = this.projectInboundMessage(runId, record); store.appendMessage(projected); this.recordInboundMessageAsThreadEvent( { @@ -514,7 +551,6 @@ export class RuntimeBridge { }, record, ); - acceptedRecords.push(record); } this.deps.db.markAgentInboxDelivered( payload.agentId, @@ -522,9 +558,6 @@ export class RuntimeBridge { ); }); } - for (const record of acceptedRecords) { - this.registerInboundInjection(runId, record); - } const priorMessages = store.loadMessages(); return { diff --git a/src/main/agent-runtime/runtime-message-injector.ts b/src/main/agent-runtime/runtime-message-injector.ts index dd72544f..60177080 100644 --- a/src/main/agent-runtime/runtime-message-injector.ts +++ b/src/main/agent-runtime/runtime-message-injector.ts @@ -27,7 +27,6 @@ import { randomUUID } from "node:crypto"; import type { AppDatabase } from "@main/database"; import type { AgentMessageRecord } from "@shared/types"; import type { ParentMessage } from "@main/worker/worker-helpers"; -import { projectInboxMessageToPi } from "./inbox-projection"; import type { RuntimeBridge } from "./runtime-bridge"; export interface RuntimeMessageInjectorDeps { @@ -90,9 +89,13 @@ export class RuntimeMessageInjector { return false; } - const projected = projectInboxMessageToPi(message, { - getAgentName: (id) => this.resolveAgentName(id), - }); + // Delegate projection + correlation registration to the bridge so + // the worker steers the exact same pi-message object that the + // bridge has indexed under `record.id`. `projectInboundMessage` + // allocates a monotonic `piTimestamp` which is what `message_end` + // surfaces back on the wire — two inbound rows that share the + // same wall-clock millisecond are no longer indistinguishable. + const projected = this.deps.runtimeBridge.projectInboundMessage(ctx.runId, message); const correlationId = randomUUID(); try { @@ -146,13 +149,6 @@ export class RuntimeMessageInjector { this.pending.delete(correlationId); if (accepted) { - try { - this.deps.runtimeBridge.registerInboundInjection(entry.runId, entry.message); - } catch (error) { - this.deps.logger.warn( - `[injector] registerInboundInjection failed for ${entry.message.id}: ${error instanceof Error ? error.message : String(error)}`, - ); - } try { this.deps.db.markAgentInboxDelivered(entry.recipientAgentId, [entry.message.id]); } catch (error) { @@ -178,13 +174,4 @@ export class RuntimeMessageInjector { } } - private resolveAgentName(agentId: string): string { - try { - const name = this.deps.db.findAgentName(agentId); - if (name && name.length > 0) return name; - } catch { - // fall through - } - return agentId; - } }