diff --git a/container/agent-runner/src/current-batch.ts b/container/agent-runner/src/current-batch.ts deleted file mode 100644 index b699c13ff..000000000 --- a/container/agent-runner/src/current-batch.ts +++ /dev/null @@ -1,29 +0,0 @@ -/** - * Per-batch context the poll loop publishes for downstream consumers - * (MCP tools, etc.) that don't sit on the poll-loop's call stack. - * - * Today the only field is `inReplyTo` — the id of the first inbound - * message in the batch the agent is currently processing. MCP tools like - * `send_message` and `send_file` read this and stamp it onto the outbound - * row so the host's a2a return-path routing can correlate replies back to - * the originating session. - * - * This is module-level state on purpose: the agent-runner is single-process - * and processes one batch at a time. Poll-loop calls `setCurrentInReplyTo` - * before invoking the provider and `clearCurrentInReplyTo` after the batch - * completes (or errors out). - */ -let currentInReplyTo: string | null = null; - -export function setCurrentInReplyTo(id: string | null): void { - currentInReplyTo = id; -} - -export function clearCurrentInReplyTo(): void { - currentInReplyTo = null; -} - -export function getCurrentInReplyTo(): string | null { - return currentInReplyTo; -} - diff --git a/container/agent-runner/src/db/session-state.ts b/container/agent-runner/src/db/session-state.ts index 9e1230926..04dc4544a 100644 --- a/container/agent-runner/src/db/session-state.ts +++ b/container/agent-runner/src/db/session-state.ts @@ -77,3 +77,47 @@ export function setContinuation(providerName: string, id: string): void { export function clearContinuation(providerName: string): void { deleteValue(continuationKey(providerName)); } + +/** + * The a2a reply stamp: the id of the first inbound message in the batch the + * agent is currently processing. The poll loop publishes it at batch start; + * MCP tools (`send_message`, `send_file`) read it and stamp it onto outbound + * rows so the host's a2a return-path routing can correlate replies back to + * the originating session. + * + * This lives in outbound.db rather than module state because the MCP server + * runs as a separate stdio subprocess from the poll loop — module state set + * by the poll loop is invisible to it. Both processes open outbound.db + * (journal_mode=DELETE + busy_timeout make intra-container access safe). + */ +const IN_REPLY_TO_KEY = 'current_in_reply_to'; + +/** + * Ignore a stamp older than this. The poll loop clears the stamp in a + * finally, but a container killed mid-batch (SIGKILL) can leave one behind; + * the guard stops a later out-of-batch read from picking up a dead stamp. + * Generous so a long-running batch's late sends still stamp correctly. + */ +const IN_REPLY_TO_MAX_AGE_MS = 30 * 60 * 1000; + +export function setCurrentInReplyTo(id: string | null): void { + if (id === null) { + clearCurrentInReplyTo(); + return; + } + setValue(IN_REPLY_TO_KEY, id); +} + +export function clearCurrentInReplyTo(): void { + deleteValue(IN_REPLY_TO_KEY); +} + +export function getCurrentInReplyTo(): string | null { + const row = getOutboundDb() + .prepare('SELECT value, updated_at FROM session_state WHERE key = ?') + .get(IN_REPLY_TO_KEY) as { value: string; updated_at: string } | undefined; + if (!row) return null; + const age = Date.now() - new Date(row.updated_at).getTime(); + if (!Number.isFinite(age) || age > IN_REPLY_TO_MAX_AGE_MS) return null; + return row.value; +} diff --git a/container/agent-runner/src/mcp-tools/core.test.ts b/container/agent-runner/src/mcp-tools/core.test.ts index 4cef950fb..8e4ff3c62 100644 --- a/container/agent-runner/src/mcp-tools/core.test.ts +++ b/container/agent-runner/src/mcp-tools/core.test.ts @@ -4,14 +4,31 @@ * batch in poll-loop, and outbound writes from MCP tools (send_message, * send_file) must pick it up so a2a return-path routing on the host can * correlate replies back to the originating session. + * + * The stamp is published through session_state in outbound.db, not module + * state — the MCP server runs as a separate stdio subprocess from the poll + * loop, so it can only see the stamp through the shared DB. These tests seed + * it the same way the poll-loop process does (a direct DB write) rather than + * via any in-memory helper, so they exercise the real process boundary. */ import { describe, it, expect, beforeEach, afterEach } from 'bun:test'; -import { initTestSessionDb, closeSessionDb, getInboundDb } from '../db/connection.js'; +import { initTestSessionDb, closeSessionDb, getInboundDb, getOutboundDb } from '../db/connection.js'; import { getUndeliveredMessages } from '../db/messages-out.js'; -import { setCurrentInReplyTo, clearCurrentInReplyTo } from '../current-batch.js'; import { sendMessage } from './core.js'; +/** + * Publish the a2a reply stamp the way the poll loop does: a direct write to + * session_state in outbound.db. `ageMs` back-dates updated_at to exercise the + * staleness guard MCP tools apply when reading it. + */ +function publishInReplyTo(id: string, ageMs = 0): void { + const updatedAt = new Date(Date.now() - ageMs).toISOString(); + getOutboundDb() + .prepare('INSERT OR REPLACE INTO session_state (key, value, updated_at) VALUES (?, ?, ?)') + .run('current_in_reply_to', id, updatedAt); +} + beforeEach(() => { initTestSessionDb(); // Seed a peer agent destination @@ -24,13 +41,12 @@ beforeEach(() => { }); afterEach(() => { - clearCurrentInReplyTo(); closeSessionDb(); }); describe('send_message MCP tool — in_reply_to plumbing', () => { - it('stamps current batch in_reply_to on outbound rows', async () => { - setCurrentInReplyTo('inbound-msg-1'); + it('stamps the batch in_reply_to (published via the DB) on outbound rows', async () => { + publishInReplyTo('inbound-msg-1'); await sendMessage.handler({ to: 'peer', text: 'hello' }); @@ -40,7 +56,17 @@ describe('send_message MCP tool — in_reply_to plumbing', () => { }); it('writes null when no batch is active', async () => { - // No setCurrentInReplyTo before this call — simulates ad-hoc / out-of-batch invocation. + // Nothing published to session_state — simulates ad-hoc / out-of-batch invocation. + await sendMessage.handler({ to: 'peer', text: 'hello' }); + + const out = getUndeliveredMessages(); + expect(out).toHaveLength(1); + expect(out[0].in_reply_to).toBeNull(); + }); + + it('ignores a stale stamp left behind by a killed container', async () => { + publishInReplyTo('inbound-msg-1', 60 * 60 * 1000); // an hour old + await sendMessage.handler({ to: 'peer', text: 'hello' }); const out = getUndeliveredMessages(); diff --git a/container/agent-runner/src/mcp-tools/core.ts b/container/agent-runner/src/mcp-tools/core.ts index 48f87d596..b9b46e087 100644 --- a/container/agent-runner/src/mcp-tools/core.ts +++ b/container/agent-runner/src/mcp-tools/core.ts @@ -9,9 +9,9 @@ import fs from 'fs'; import path from 'path'; -import { getCurrentInReplyTo } from '../current-batch.js'; import { findByName, getAllDestinations } from '../destinations.js'; import { getMessageIdBySeq, getRoutingBySeq, writeMessageOut } from '../db/messages-out.js'; +import { getCurrentInReplyTo } from '../db/session-state.js'; import { getSessionRouting } from '../db/session-routing.js'; import { registerTools } from './server.js'; import type { McpToolDefinition } from './types.js'; diff --git a/container/agent-runner/src/poll-loop.ts b/container/agent-runner/src/poll-loop.ts index c46e67aab..19e9c09f6 100644 --- a/container/agent-runner/src/poll-loop.ts +++ b/container/agent-runner/src/poll-loop.ts @@ -2,8 +2,13 @@ import { findByName, getAllDestinations, type DestinationEntry } from './destina import { getPendingMessages, markProcessing, markCompleted, type MessageInRow } from './db/messages-in.js'; import { writeMessageOut } from './db/messages-out.js'; import { getInboundDb, touchHeartbeat, clearStaleProcessingAcks } from './db/connection.js'; -import { clearContinuation, migrateLegacyContinuation, setContinuation } from './db/session-state.js'; -import { clearCurrentInReplyTo, setCurrentInReplyTo } from './current-batch.js'; +import { + clearContinuation, + clearCurrentInReplyTo, + migrateLegacyContinuation, + setContinuation, + setCurrentInReplyTo, +} from './db/session-state.js'; import { formatMessages, extractRouting,