Gate inbox delivery on worker-side steer ACK

Replace the fire-and-forget steer with a correlated ACK round-trip so main
only marks an inbox row delivered once the worker confirms the message
landed on the live Agent. Steer against a run that is still in the
`preparing` window or has already started `closing` now falls back to the
wake + inbox projection path instead of being silently dropped.

Worker publishes `run-agent-state` transitions at Agent construction and
teardown; main tracks these on `ActiveRunContext.state` and refuses to
steer runs that are not `agent_active`. A rejected or timed-out ACK
triggers `wakeAgentIfPossible` so the next `prepareRun` replays the row.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-19 14:38:04 -04:00
1 parent ad9205fcc9
commit e917c614c9
12 files changed
+264 -57

No files matched your search

+19
View File
@@ -30,6 +30,24 @@ import {
} from "./artifact-projector";
import { compactMessagesIfNeeded } from "./context-compactor";
/**
* Lifecycle phase of an active pi run as seen by main.
*
* - `preparing`: `prepareRun` RPC has landed; main knows the run is real,
* but the worker has not yet constructed the Agent instance. The
* in-process `activeAgents` map is empty, so any `steer` attempt would
* be silently dropped.
* - `agent_active`: the worker has both constructed the Agent and
* subscribed to its event stream. Steer-run messages are safe to
* dispatch — `agent.steer()` will enqueue onto the live steering queue
* and pi will drain it at the next turn boundary.
* - `closing`: the Agent finished (or aborted) and the worker has
* unsubscribed and removed the instance from `activeAgents`. The
* `run-finished` child message may still be in flight; treat steer
* messages as rejected so inbox delivery falls back to the next run.
*/
export type ActiveRunLifecycleState = "preparing" | "agent_active" | "closing";
export interface ActiveRunContext {
runId: string;
agentId: string;
@@ -39,6 +57,7 @@ export interface ActiveRunContext {
modelId: string;
taskSessionId: string;
workspacePath: string | null;
state: ActiveRunLifecycleState;
}
export interface HookHandlerDeps {
+25 -3
View File
@@ -13,6 +13,7 @@ import {
handleAfterToolCall,
handleTransformContext,
type ActiveRunContext,
type ActiveRunLifecycleState,
} from "./hook-handlers";
import type {
AfterToolCallRpcPayload,
@@ -79,14 +80,34 @@ export class RuntimeBridge {
/** Returns the active run context for a given agent, if any. Used by
* main-side services (inbox dispatcher) to decide whether a new inbox
* row should route to `agent.steer()` instead of waiting for the next
* `prepareRun` cycle. */
findActiveRunForAgent(agentId: string): ActiveRunContext | null {
* `prepareRun` cycle. When `requireSteerReady` is set, only runs whose
* worker-side Agent is actually subscribed (`agent_active`) are
* returned — this avoids steering into `preparing` or `closing` runs
* where the steer would be silently dropped. */
findActiveRunForAgent(
agentId: string,
options?: { requireSteerReady?: boolean },
): ActiveRunContext | null {
for (const ctx of this.activeRuns.values()) {
if (ctx.agentId === agentId) return ctx;
if (ctx.agentId !== agentId) continue;
if (options?.requireSteerReady && ctx.state !== "agent_active") continue;
return ctx;
}
return null;
}
/**
* Transition an active run between lifecycle phases as the worker
* publishes `run-agent-state`. Unknown run ids are ignored: the worker
* may publish state for a run that main has already evicted when a
* `run-finished` races ahead of the trailing state message.
*/
setRunAgentState(runId: string, state: ActiveRunLifecycleState): void {
const ctx = this.activeRuns.get(runId);
if (!ctx) return;
ctx.state = state;
}
private getOrCreateStore(taskSessionId: string): SqliteSessionStore {
const existing = this.sessionStores.get(taskSessionId);
if (existing) return existing;
@@ -268,6 +289,7 @@ export class RuntimeBridge {
modelId: payload.modelId,
taskSessionId: payload.taskSessionId,
workspacePath: payload.workspacePath,
state: "preparing",
});
const definition = getConnectorDefinition(payload.connectorId);
const resolved = await this.deps.credentialResolver.resolve(payload.connectorId);
@@ -1,53 +1,86 @@
/**
* Runtime message injector — the inward half of the Slack-like
* messaging system. When an agent has an active pi run, incoming
* direct / channel / mention messages must reach the agent *during*
* the loop, not on the next heartbeat. Pre-Task-3 the only delivery
* path was a batched `prepareRun` projection, so any message that
* arrived while the recipient was working sat dark until the next
* wake.
* Runtime message injector — the inward half of the Slack-like messaging
* system. When an agent has an active pi run, incoming direct, channel,
* or mention messages must reach the agent during the loop rather than
* on the next heartbeat. Pre-injector, the only delivery path was the
* batched `prepareRun` projection, so any message that arrived while the
* recipient was running sat dark until the next wake.
*
* This module closes that gap by routing a projected pi AgentMessage
* into the worker's active Agent via the `steer-run` ParentMessage.
* Pi's steering queue drains at the next turn boundary; the agent
* observes the new user turn and can react in-loop.
* Delivery is a two-step handshake:
*
* Delivery is best-effort and marked as delivered upon post. The
* session-store shadow still writes via the normal pi message_end
* event so the transcript remains durable. Callers fall back to the
* existing inbox+wake path when `tryInject` returns false.
* main → worker steer-run { correlationId, runId, message }
* worker → main steer-ack { correlationId, accepted }
*
* The worker's `runner.steer()` is the single authority for whether the
* message reached the live Agent. Main holds the inbox `delivered_at`
* write until the ACK lands: a positive ACK commits the write; a rejected
* or timed-out ACK leaves the row unread and re-triggers the recipient's
* wake, so the next `prepareRun` picks the message up via the regular
* inbox projection. Either way the message is observed exactly once.
*
* The injector short-circuits when the active run is not in
* `agent_active` — `preparing` and `closing` windows cannot accept a
* steer — and the caller falls back to the wake + inbox flow directly.
*/
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";
import type { ParentMessage } from "@main/worker/worker-helpers";
export interface RuntimeMessageInjectorDeps {
db: AppDatabase;
runtimeBridge: RuntimeBridge;
postToWorker: (message: ParentMessage) => void;
/**
* Lazy wake hook. The injector is instantiated before orchestration so
* the wake function has to be resolved per-call; returning `null` means
* orchestration is not yet online, in which case the caller will have
* handled the wake path already.
*/
getWakeAgent: () => ((agentId: string, companyId: string, trigger: string) => string | null) | null;
logger: {
info: (...args: unknown[]) => void;
warn: (...args: unknown[]) => void;
};
steerAckTimeoutMs?: number;
}
interface PendingSteer {
recipientAgentId: string;
companyId: string;
messageId: string;
timer: ReturnType<typeof setTimeout>;
}
const DEFAULT_STEER_ACK_TIMEOUT_MS = 5_000;
export class RuntimeMessageInjector {
constructor(private readonly deps: RuntimeMessageInjectorDeps) {}
private readonly pending = new Map<string, PendingSteer>();
private readonly ackTimeoutMs: number;
constructor(private readonly deps: RuntimeMessageInjectorDeps) {
this.ackTimeoutMs = deps.steerAckTimeoutMs ?? DEFAULT_STEER_ACK_TIMEOUT_MS;
}
/**
* Attempts to deliver a persisted inbox message into the recipient's
* live pi run. Returns true when the message was dispatched to the
* worker, false when no active run exists for the recipient.
* live pi run. Returns true when the steer was dispatched to the worker
* with a pending ACK — callers must skip the synchronous wake fallback
* for this recipient because the injector now owns the message's fate:
* it will either mark the row delivered on a positive ACK or trigger a
* wake on a negative/timed-out ACK.
*
* The caller owns the "post-delivery" decision: on true, it should
* skip the regular `wakeAgentIfPossible` path for this recipient; on
* false, it falls back to the legacy inbox+wake flow.
* Returns false when no steer was dispatched (no active run, cross-
* company mismatch, or the run is not yet `agent_active`). Callers then
* fall through to the legacy inbox + wake flow immediately.
*/
tryInject(recipientAgentId: string, companyId: string, message: AgentMessageRecord): boolean {
const ctx = this.deps.runtimeBridge.findActiveRunForAgent(recipientAgentId);
const ctx = this.deps.runtimeBridge.findActiveRunForAgent(recipientAgentId, {
requireSteerReady: true,
});
if (!ctx) return false;
if (ctx.companyId !== companyId) {
this.deps.logger.warn(
@@ -59,9 +92,15 @@ export class RuntimeMessageInjector {
const projected = projectInboxMessageToPi(message, {
getAgentName: (id) => this.resolveAgentName(id),
});
const correlationId = randomUUID();
try {
this.deps.postToWorker({ type: "steer-run", runId: ctx.runId, message: projected });
this.deps.postToWorker({
type: "steer-run",
runId: ctx.runId,
correlationId,
message: projected,
});
} catch (error) {
this.deps.logger.warn(
`[injector] postToWorker steer-run failed for run ${ctx.runId}: ${error instanceof Error ? error.message : String(error)}`,
@@ -69,14 +108,65 @@ export class RuntimeMessageInjector {
return false;
}
const timer = setTimeout(() => {
this.resolvePending(correlationId, false, "timeout");
}, this.ackTimeoutMs);
if (typeof timer.unref === "function") timer.unref();
this.pending.set(correlationId, {
recipientAgentId,
companyId,
messageId: message.id,
timer,
});
return true;
}
/**
* Routed from `worker-message-handler` on every `steer-ack`. Commits
* the inbox `delivered_at` write when the worker confirms the steer
* reached the live Agent; otherwise keeps the row unread and asks
* orchestration to wake the recipient so the next `prepareRun` picks
* the message up through the regular inbox projection.
*/
handleSteerAck(correlationId: string, accepted: boolean): void {
this.resolvePending(correlationId, accepted, "worker_ack");
}
private resolvePending(
correlationId: string,
accepted: boolean,
source: "worker_ack" | "timeout",
): void {
const entry = this.pending.get(correlationId);
if (!entry) return;
clearTimeout(entry.timer);
this.pending.delete(correlationId);
if (accepted) {
try {
this.deps.db.markAgentInboxDelivered(entry.recipientAgentId, [entry.messageId]);
} catch (error) {
this.deps.logger.warn(
`[injector] markAgentInboxDelivered failed for ${entry.messageId}: ${error instanceof Error ? error.message : String(error)}`,
);
}
return;
}
this.deps.logger.warn(
`[injector] steer rejected (${source}) for message ${entry.messageId} to agent ${entry.recipientAgentId}; falling back to wake + inbox projection`,
);
const wake = this.deps.getWakeAgent();
if (!wake) return;
try {
this.deps.db.markAgentInboxDelivered(recipientAgentId, [message.id]);
wake(entry.recipientAgentId, entry.companyId, "message");
} catch (error) {
this.deps.logger.warn(
`[injector] markAgentInboxDelivered failed for ${message.id}: ${error instanceof Error ? error.message : String(error)}`,
`[injector] wake fallback failed for agent ${entry.recipientAgentId}: ${error instanceof Error ? error.message : String(error)}`,
);
}
return true;
}
private resolveAgentName(agentId: string): string {
+6
View File
@@ -171,6 +171,7 @@ export function createServiceGraph(deps: ServiceGraphDependencies) {
budgetService: leaf.budgetEnforcementService,
sessionService: leaf.sessionManagementService,
piRuntimeRef,
getRuntimeMessageInjector: () => runtimeMessageInjector,
});
runtimeMessageInjector = new RuntimeMessageInjector({
@@ -181,6 +182,11 @@ export function createServiceGraph(deps: ServiceGraphDependencies) {
if (!worker) throw new Error("Worker is not initialised — cannot deliver steer message.");
worker.postMessage(message);
},
getWakeAgent: () => {
if (!orchestrationRef.isAssigned()) return null;
const orch = orchestrationRef.get();
return (agentId, companyId, trigger) => orch.wakeAgentIfPossible(agentId, companyId, trigger);
},
logger,
});
@@ -7,6 +7,7 @@ import type { createBudgetEnforcementService } from "../budget-enforcement-servi
import type { createSessionManagementService } from "../session-management-service";
import type { PiRuntimeServices } from "../service-graph-pi-runtime";
import type { NotificationOptions } from "../notification-service";
import type { RuntimeMessageInjector } from "../agent-runtime";
import type { LateBinding } from "./late-binding";
export interface WorkerServiceDeps {
@@ -29,6 +30,12 @@ export interface WorkerServiceDeps {
budgetService: ReturnType<typeof createBudgetEnforcementService>;
sessionService: ReturnType<typeof createSessionManagementService>;
piRuntimeRef: { current: PiRuntimeServices | null };
/**
* Resolver for the late-constructed runtime message injector. The worker
* message handler needs it to route `steer-ack` child messages, but the
* injector is assembled after the worker service itself.
*/
getRuntimeMessageInjector: () => RuntimeMessageInjector | null;
}
/**
@@ -52,6 +59,7 @@ export function buildWorkerService(deps: WorkerServiceDeps) {
budgetService,
sessionService,
piRuntimeRef,
getRuntimeMessageInjector,
} = deps;
const wakeAgentIfPossible = (agentId: string | null | undefined, companyId: string, trigger: string) =>
@@ -83,5 +91,6 @@ export function buildWorkerService(deps: WorkerServiceDeps) {
budgetService,
sessionService,
getRuntimeBridge: () => piRuntimeRef.current?.runtimeBridge ?? null,
getRuntimeMessageInjector,
});
}
+4
View File
@@ -26,6 +26,10 @@ export class LateBinding<T> {
return this.value;
}
isAssigned(): boolean {
return this.assigned;
}
clear(): void {
this.value = null;
this.assigned = false;
+16 -1
View File
@@ -21,6 +21,7 @@ export function createWorkerService({
budgetService,
sessionService,
getRuntimeBridge,
getRuntimeMessageInjector,
}: WorkerServiceDependencies) {
const MAX_RETRY_ATTEMPTS = 2;
const RETRY_COOLDOWN_MS = 30_000; // 30 seconds
@@ -111,7 +112,21 @@ export function createWorkerService({
await handleWorkerMessage(
message,
worker,
{ db, logger, isQuitting, emitEvent, wakeAgentIfPossible, processRunCompletion, onRunStarted, onRunCompleted, notify, budgetService, sessionService, getRuntimeBridge },
{
db,
logger,
isQuitting,
emitEvent,
wakeAgentIfPossible,
processRunCompletion,
onRunStarted,
onRunCompleted,
notify,
budgetService,
sessionService,
getRuntimeBridge,
getRuntimeMessageInjector,
},
workerMessageSchema,
destroyed,
MAX_RETRY_ATTEMPTS,
+7 -5
View File
@@ -129,11 +129,13 @@ parentPort.on("message", (event) => {
return;
}
if (message.type === "steer-run") {
runner.steer(message.runId, message.message);
return;
}
if (message.type === "followup-run") {
runner.followUp(message.runId, message.message);
const accepted = runner.steer(message.runId, message.message);
post({
type: "steer-ack",
runId: message.runId,
correlationId: message.correlationId,
accepted,
});
return;
}
if (message.type === "rpc-response") {
+16 -7
View File
@@ -24,17 +24,24 @@ export interface RunJob {
}
/**
* Parent → worker control-plane messages. `steer-run` and `followup-run`
* are the inward half of pi's mid-run injection contract: the worker owns
* the Agent instance, main owns the inbox DB, so main posts the enqueued
* user message here and the worker calls `agent.steer()` / `agent.followUp()`.
* Parent → worker control-plane messages. `steer-run` is the inward half
* of pi's mid-run injection contract: the worker owns the Agent instance,
* main owns the inbox DB, so main posts the enqueued user message here and
* the worker calls `agent.steer()`.
*
* `correlationId` pairs each steer with a `steer-ack` child reply so main
* can hold the "inbox delivered" write until the worker has queued the
* message onto the live Agent. Without the ACK, a steer that lost the
* race against run teardown would be silently dropped after main already
* marked the row delivered.
*/
export type RunAgentLifecycleState = "agent_active" | "closing";
export type ParentMessage =
| { type: "enqueue-run"; payload: RunJob }
| { type: "cancel-run"; runId: string }
| { type: "shutdown" }
| { type: "steer-run"; runId: string; message: AgentMessage }
| { type: "followup-run"; runId: string; message: AgentMessage }
| { type: "steer-run"; runId: string; correlationId: string; message: AgentMessage }
| { type: "rpc-response"; rpcId: string; response: { ok: true; result: unknown } | { ok: false; error: string } };
export type ChildMessage =
@@ -62,7 +69,9 @@ export type ChildMessage =
emittedAt: string;
event: { type: string } & Record<string, unknown>;
}
| { type: "rpc-request"; rpcId: string; runId: string; method: string; payload: unknown };
| { type: "rpc-request"; rpcId: string; runId: string; method: string; payload: unknown }
| { type: "run-agent-state"; runId: string; state: RunAgentLifecycleState }
| { type: "steer-ack"; runId: string; correlationId: string; accepted: boolean };
export function buildRedactor(env: Record<string, string>): (text: string) => string {
const sensitivePatterns: string[] = [];
+13 -2
View File
@@ -1,4 +1,4 @@
import type { RuntimeBridge } from "../agent-runtime";
import type { RuntimeBridge, RuntimeMessageInjector } from "../agent-runtime";
import type { workerMessageSchema } from "./worker-types";
import { handleRunLog } from "./handlers/run-log-handler";
import { handleRunStatus } from "./handlers/run-status-handler";
@@ -8,7 +8,10 @@ import type { HandlerContext, HandlerDeps } from "./handlers/handler-context";
export async function handleWorkerMessage(
message: unknown,
worker: ReturnType<typeof import("electron").utilityProcess.fork> | null,
deps: HandlerDeps & { getRuntimeBridge?: () => RuntimeBridge | null },
deps: HandlerDeps & {
getRuntimeBridge?: () => RuntimeBridge | null;
getRuntimeMessageInjector?: () => RuntimeMessageInjector | null;
},
workerMessageSchemaRef: typeof workerMessageSchema,
destroyed: boolean,
MAX_RETRY_ATTEMPTS: number,
@@ -83,5 +86,13 @@ export async function handleWorkerMessage(
}
return;
}
case "run-agent-state": {
deps.getRuntimeBridge?.()?.setRunAgentState(event.runId, event.state);
return;
}
case "steer-ack": {
deps.getRuntimeMessageInjector?.()?.handleSteerAck(event.correlationId, event.accepted);
return;
}
}
}
+15 -11
View File
@@ -28,8 +28,10 @@ interface RunnerDeps {
* - Builds an `Agent` via `createPiAgent` with the full Pi contract.
* - Streams every pi `AgentEvent` into the worker → main pipe.
* - Persists each `message_end` message through `persistMessage` RPC.
* - Exposes `steer` / `followUp` to the worker entry so mid-run
* injections from main land on the in-process Agent instance.
* - Exposes `steer` to the worker entry so mid-run injections from
* main land on the in-process Agent instance.
* - Publishes `run-agent-state` transitions so main can tell whether
* an active run is actually ready to accept steer messages.
*/
export class WorkerRunner {
private readonly abortControllers = new Map<string, AbortController>();
@@ -43,7 +45,15 @@ export class WorkerRunner {
if (controller) controller.abort();
}
/** Enqueue a steering message onto the active Agent for this run. */
/**
* Enqueue a steering message onto the active Agent for this run. Returns
* true iff the Agent instance exists in `activeAgents` and the message
* was appended to its steering queue. A false return means the run has
* either not finished construction (pre-`agent_active`) or has already
* torn down (post-`closing`); callers must treat the steer as rejected
* so main can hold the inbox `delivered_at` write and fall back to the
* inbox + wake path instead.
*/
steer(runId: string, message: AgentMessage): boolean {
const agent = this.activeAgents.get(runId);
if (!agent) return false;
@@ -51,14 +61,6 @@ export class WorkerRunner {
return true;
}
/** Enqueue a follow-up message onto the active Agent for this run. */
followUp(runId: string, message: AgentMessage): boolean {
const agent = this.activeAgents.get(runId);
if (!agent) return false;
agent.followUp(message);
return true;
}
/**
* Aborts every in-flight run and returns a promise that resolves once
* each `run()` has settled (success, failure, or abort). Callers (the
@@ -150,6 +152,7 @@ export class WorkerRunner {
rpc,
});
this.activeAgents.set(job.runId, agent);
post({ type: "run-agent-state", runId: job.runId, state: "agent_active" });
const unsubscribe = agent.subscribe(async (event) => {
const timestamp = new Date().toISOString();
@@ -192,6 +195,7 @@ export class WorkerRunner {
abortController.signal.addEventListener("abort", () => agent.abort(), { once: true });
await runPromise;
} finally {
post({ type: "run-agent-state", runId: job.runId, state: "closing" });
unsubscribe();
this.activeAgents.delete(job.runId);
}
+17 -1
View File
@@ -4,7 +4,7 @@ import type { DesktopEvent, RunRecord } from "@shared/types";
import type { NotificationOptions } from "../notification-service";
import type { createBudgetEnforcementService } from "../budget-enforcement-service";
import type { createSessionManagementService } from "../session-management-service";
import type { RuntimeBridge } from "../agent-runtime";
import type { RuntimeBridge, RuntimeMessageInjector } from "../agent-runtime";
const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]);
@@ -56,6 +56,17 @@ export const workerMessageSchema = z.union([
]),
payload: z.unknown(),
}),
z.object({
type: z.literal("run-agent-state"),
runId: z.string().uuid(),
state: z.enum(["agent_active", "closing"]),
}),
z.object({
type: z.literal("steer-ack"),
runId: z.string().uuid(),
correlationId: z.string(),
accepted: z.boolean(),
}),
]);
export interface WorkerServiceDependencies {
@@ -92,4 +103,9 @@ export interface WorkerServiceDependencies {
sessionService?: ReturnType<typeof createSessionManagementService>;
/** Forwards `agent-event` from worker to renderer and routes RPC to main services. */
getRuntimeBridge?: () => RuntimeBridge | null;
/**
* Resolves the shared runtime message injector so `steer-ack` child
* messages can commit inbox delivery or trigger a wake fallback.
*/
getRuntimeMessageInjector?: () => RuntimeMessageInjector | null;
}