diff --git a/src/main/services/agent-run-loop.ts b/src/main/services/agent-run-loop.ts index 2cc87f82..02a8a44d 100644 --- a/src/main/services/agent-run-loop.ts +++ b/src/main/services/agent-run-loop.ts @@ -17,6 +17,8 @@ export interface AgentRunLoopInput { initialUserTurn: string; toolContext: AgentToolContext; externalTools?: ToolDefinition[]; + /** Messages that arrived mid-run, or null. */ + pollNewMessages?: () => string | null; } export interface AgentRunLoopResult { @@ -63,6 +65,11 @@ export async function runAgentLoop(deps: AgentRunLoopDeps, input: AgentRunLoopIn // completing the run would silently mark a no-op as done. throw new Error("Provider returned an empty turn: no text and no tool calls."); } + const lateMessages = input.pollNewMessages?.(); + if (lateMessages) { + messages.push({ role: "assistant", text: turn.text, toolCalls: [] }, { role: "user", text: lateMessages }); + continue; + } return { finalText: turn.text, iterations: iteration, toolCallCount, lastRequestId }; } @@ -95,7 +102,8 @@ export async function runAgentLoop(deps: AgentRunLoopDeps, input: AgentRunLoopIn } results.push({ toolCallId: call.id, toolName: call.name, output }); } - messages.push({ role: "tool", results }); + const followUp = input.pollNewMessages?.(); + messages.push(followUp ? { role: "tool", results, followUp } : { role: "tool", results }); deps.onToolResults?.(); } diff --git a/src/main/services/chat-wake-context.ts b/src/main/services/chat-wake-context.ts index 18978ba7..46639670 100644 --- a/src/main/services/chat-wake-context.ts +++ b/src/main/services/chat-wake-context.ts @@ -52,3 +52,32 @@ function formatMessage(m: ThreadMessageRow): string { function utcMinute(iso: string): string { return `${iso.slice(0, 10)} ${iso.slice(11, 16)} UTC`; } + +/** Polls the threads this run is part of for messages from others. */ +export function createRunMessagePoller( + db: DatabaseSync, + input: { orgId: string; agentId: string; sourceThreadId: string | null; runStartedAt: string }, +): () => string | null { + let watermark = input.runStartedAt; + return () => { + const rows = db.prepare(` + select m.body_md, m.created_at, m.edited_at, m.sender_type, m.sender_id, + coalesce(a.name, e.display_name, m.sender_type) as sender_name, + coalesce(a.role, e.role) as sender_role + from chat_messages m + left join agents a on a.id = m.sender_id and m.sender_type = 'agent' + left join employees e on e.id = m.sender_id and m.sender_type = 'employee' + where m.org_id = ? and m.created_at > ? and m.deleted_at is null + and not (m.sender_type = 'agent' and m.sender_id = ?) + and (m.thread_id = ? or m.thread_id in ( + select own.thread_id from chat_messages own + where own.org_id = ? and own.sender_type = 'agent' and own.sender_id = ? and own.created_at >= ? + )) + order by m.created_at asc + `).all(input.orgId, watermark, input.agentId, input.sourceThreadId, input.orgId, input.agentId, input.runStartedAt) as unknown as ThreadMessageRow[]; + if (rows.length === 0) return null; + watermark = rows[rows.length - 1].created_at; + console.info(`[run-inbox] agent ${input.agentId}: ${rows.length} new message(s) injected mid-run`); + return ["Messages that arrived while you were working (they may change what you should do):", ...rows.map(formatMessage)].join("\n"); + }; +} diff --git a/src/main/services/provider-tool-turn.ts b/src/main/services/provider-tool-turn.ts index 867db307..970957de 100644 --- a/src/main/services/provider-tool-turn.ts +++ b/src/main/services/provider-tool-turn.ts @@ -24,7 +24,7 @@ export interface ToolCall { export type AgentTurnMessage = | { role: "user"; text: string } | { role: "assistant"; text: string; toolCalls: ToolCall[] } - | { role: "tool"; results: Array<{ toolCallId: string; toolName: string; output: string }> }; + | { role: "tool"; results: Array<{ toolCallId: string; toolName: string; output: string }>; followUp?: string }; export interface AgentTurnResult { text: string; diff --git a/src/main/services/provider-tool-turns/anthropic-format.ts b/src/main/services/provider-tool-turns/anthropic-format.ts index ed98ba72..fdcf3b2f 100644 --- a/src/main/services/provider-tool-turns/anthropic-format.ts +++ b/src/main/services/provider-tool-turns/anthropic-format.ts @@ -17,7 +17,9 @@ export function toAnthropicMessages(messages: AgentTurnMessage[]): unknown[] { for (const call of message.toolCalls) content.push({ type: "tool_use", id: call.id, name: call.name, input: call.arguments }); out.push({ role: "assistant", content }); } else { - out.push({ role: "user", content: message.results.map((result) => ({ type: "tool_result", tool_use_id: result.toolCallId, content: result.output })) }); + const content: unknown[] = message.results.map((result) => ({ type: "tool_result", tool_use_id: result.toolCallId, content: result.output })); + if (message.followUp) content.push({ type: "text", text: message.followUp }); + out.push({ role: "user", content }); } } return out; diff --git a/src/main/services/provider-tool-turns/bedrock-format.ts b/src/main/services/provider-tool-turns/bedrock-format.ts index 3a14319e..3be38e99 100644 --- a/src/main/services/provider-tool-turns/bedrock-format.ts +++ b/src/main/services/provider-tool-turns/bedrock-format.ts @@ -20,7 +20,9 @@ export function toBedrockMessages(messages: AgentTurnMessage[]): JsonRecord[] { for (const call of message.toolCalls) content.push({ toolUse: { toolUseId: call.id, name: call.name, input: call.arguments } }); out.push({ role: "assistant", content }); } else { - out.push({ role: "user", content: message.results.map((result) => ({ toolResult: { toolUseId: result.toolCallId, content: [{ text: result.output }] } })) }); + const content: unknown[] = message.results.map((result) => ({ toolResult: { toolUseId: result.toolCallId, content: [{ text: result.output }] } })); + if (message.followUp) content.push({ text: message.followUp }); + out.push({ role: "user", content }); } } return out; diff --git a/src/main/services/provider-tool-turns/chat-format.ts b/src/main/services/provider-tool-turns/chat-format.ts index 2dcf05f4..02e87bcf 100644 --- a/src/main/services/provider-tool-turns/chat-format.ts +++ b/src/main/services/provider-tool-turns/chat-format.ts @@ -22,6 +22,7 @@ export function toChatMessages(instructions: string | null, messages: AgentTurnM for (const result of message.results) { out.push({ role: "tool", tool_call_id: result.toolCallId, content: result.output }); } + if (message.followUp) out.push({ role: "user", content: message.followUp }); } } return out; diff --git a/src/main/services/provider-tool-turns/gemini-format.ts b/src/main/services/provider-tool-turns/gemini-format.ts index 1bc7424e..bae9c5a5 100644 --- a/src/main/services/provider-tool-turns/gemini-format.ts +++ b/src/main/services/provider-tool-turns/gemini-format.ts @@ -43,9 +43,10 @@ function toGeminiContents(model: string, messages: AgentTurnMessage[]): JsonReco } if (parts.length > 0) contents.push({ role: "model", parts }); } else { - const responses = message.results.map((result) => ({ + const responses: unknown[] = message.results.map((result) => ({ functionResponse: { name: result.toolName, response: { output: result.output }, ...(withId ? { id: result.toolCallId } : {}) }, })); + if (message.followUp) responses.push({ text: message.followUp }); const last = contents[contents.length - 1]; if (last && last.role === "user" && Array.isArray(last.parts) && last.parts.some((part) => isRecord(part) && "functionResponse" in part)) { (last.parts as unknown[]).push(...responses); diff --git a/src/main/services/provider-tool-turns/responses-format.ts b/src/main/services/provider-tool-turns/responses-format.ts index 51f023a8..e1f469c0 100644 --- a/src/main/services/provider-tool-turns/responses-format.ts +++ b/src/main/services/provider-tool-turns/responses-format.ts @@ -35,6 +35,7 @@ export function toResponsesInput(messages: AgentTurnMessage[]): unknown[] { for (const result of message.results) { items.push({ type: "function_call_output", call_id: result.toolCallId, output: result.output }); } + if (message.followUp) items.push({ role: "user", content: [{ type: "input_text", text: message.followUp }] }); } } return items; diff --git a/src/main/services/task-execution-service.ts b/src/main/services/task-execution-service.ts index c016545d..8ef3d5af 100644 --- a/src/main/services/task-execution-service.ts +++ b/src/main/services/task-execution-service.ts @@ -24,7 +24,7 @@ import { type WorkspaceRow, } from "./task-prompts"; import { renderTaskHandoff } from "./agent-task-handoff"; -import { loadChatWakeContext } from "./chat-wake-context"; +import { createRunMessagePoller, loadChatWakeContext } from "./chat-wake-context"; import { loadComposioTools } from "./composio/composio-tool-bridge"; import { connectorDisplayName } from "./composio/connector-brands"; import { ConnectorGrantsService } from "./composio/connector-grants"; @@ -134,6 +134,8 @@ export class TaskExecutionService { let skillRuntimeSnapshot: SkillRuntimeSnapshotRow | null = null; try { skillRuntimeSnapshot = this.resolveSkillRuntimeSnapshot(orgId, agent.id, runId); + // Watermark for reply delivery and mid-run message polling. + const runStartedAt = new Date().toISOString(); const chatContext = loadChatWakeContext(this.db, orgId, agent.id, isolation.secretScopeJson); const toolContext: AgentToolContext = { orgId, @@ -153,9 +155,6 @@ export class TaskExecutionService { .map(connectorDisplayName); const language = this.appLanguage(); const externalTools = await loadComposioTools(this.db, this.providers.integrationKeys, orgId, agent.id); - // Watermark for the reply-delivery check below: anything the agent posts - // into the source thread from here on counts as its reply. - const runStartedAt = new Date().toISOString(); const result = await runAgentLoop( { providers: this.providers, @@ -185,6 +184,7 @@ export class TaskExecutionService { }), toolContext, externalTools, + pollNewMessages: createRunMessagePoller(this.db, { orgId, agentId: agent.id, sourceThreadId: chatContext?.threadId ?? null, runStartedAt }), }, ); providerRequestId = result.lastRequestId; diff --git a/tests/integration/chat-context.test.ts b/tests/integration/chat-context.test.ts index 446d73a9..b1826504 100644 --- a/tests/integration/chat-context.test.ts +++ b/tests/integration/chat-context.test.ts @@ -19,10 +19,15 @@ const testVault = { async function founderLine() { const wakePrompts: string[] = []; + const wakeInputs: ProviderToolTurnInput["messages"][] = []; + let onWake: (turnIndex: number) => void = () => undefined; const turn = async (input: ProviderToolTurnInput): Promise => { const first = input.messages.find((message) => message.role === "user"); const prompt = first && first.role === "user" ? first.text : ""; - if (/thread_id="/.test(prompt) && input.messages.length === 1) wakePrompts.push(prompt); + if (!/thread_id="/.test(prompt)) return { text: "Kickoff noted.", toolCalls: [], rawResponseJson: "{}" }; + if (input.messages.length === 1) wakePrompts.push(prompt); + wakeInputs.push(input.messages); + onWake(wakeInputs.length); return { text: `Reply ${wakePrompts.length}.`, toolCalls: [], rawResponseJson: "{}" }; }; const connection = new DatabaseConnection(join(tmpdir(), `ac-chatctx-${randomUUID()}.db`)); @@ -60,7 +65,8 @@ async function founderLine() { const row = connection.db.prepare("select id from chat_messages where body_md = ?").get(bodyMd) as { id: string }; chat.editMessage({ orgId, messageId: row.id, bodyMd: newBody, actorType: "employee", actorId: founderId }); }; - return { connection, wakePrompts, say, runWake, edit }; + const onWakeTurn = (hook: (turnIndex: number) => void) => { onWake = hook; }; + return { connection, wakePrompts, wakeInputs, say, runWake, edit, onWakeTurn }; } describe("chat wake context", () => { @@ -99,4 +105,72 @@ describe("chat wake context", () => { expect(fresh).not.toContain("Reply 1."); h.connection.close(); }); + + test("a message that arrives mid-run reaches the next model turn", async () => { + const connection = new DatabaseConnection(join(tmpdir(), `ac-midrun-${randomUUID()}.db`)); + const chat = new ChatService(connection.db); + const seenFollowUps: string[] = []; + let founderId = ""; + let orgId = ""; + let spaceId = ""; + let threadId = ""; + let wakeTurns = 0; + const turn = async (input: ProviderToolTurnInput): Promise => { + const first = input.messages.find((message) => message.role === "user"); + if (!(first && first.role === "user" && /thread_id="/.test(first.text))) return { text: "Kickoff noted.", toolCalls: [], rawResponseJson: "{}" }; + wakeTurns += 1; + const last = input.messages[input.messages.length - 1]; + if (last.role === "tool" && last.followUp) seenFollowUps.push(last.followUp); + if (wakeTurns === 1) { + chat.sendMessage({ orgId, spaceId, threadId, senderType: "employee", senderId: founderId, bodyMd: "Stop — use Postgres instead.", attachmentFileIds: [] }); + return { text: "Looking.", toolCalls: [{ id: "c1", name: "read_messages", arguments: {} }], rawResponseJson: "{}" }; + } + return { text: "Switching to Postgres.", toolCalls: [], rawResponseJson: "{}" }; + }; + const providers = new ProviderService(connection.db, providerAdapters, testVault, undefined, turn); + const onboarding = new OnboardingService({ db: connection.db, saveProviderCredential: (input) => providers.saveCredential(input) }); + const execution = new TaskExecutionService(connection.db, providers); + let ceoId = ""; + connection.transaction(() => { + const summary = onboarding.bootstrap({ + companyName: "FlowPilot", companyDescription: "CSM analytics.", employeeName: "Yang", workspacePath: "", + executionMode: "ai_ceo", autoApproveHires: true, reviewDeliverables: false, + goalTitle: "Ship the MVP.", goalDescription: "A working product.", + providerKey: "openai", providerSecret: "sk-test-key", providerModel: "gpt-5.4", + }); + orgId = summary.orgId; + ceoId = summary.ceoAgentId; + }); + founderId = (connection.db.prepare("select id from employees where org_id = ?").get(orgId) as { id: string }).id; + spaceId = ensureCeoDirectSpace(connection.db, orgId)!; + const kickoff = claimNextQueuedRun(connection); + await execution.executeRun(kickoff!.org_id, kickoff!.id); + const firstId = chat.sendMessage({ orgId, spaceId, threadId: null, senderType: "employee", senderId: founderId, bodyMd: "Use MySQL.", attachmentFileIds: [] }); + threadId = (connection.db.prepare("select thread_id from chat_messages where id = ?").get(firstId) as { thread_id: string }).thread_id; + const wake = connection.db.prepare( + "select id from task_runs where org_id = ? and agent_id = ? and trigger = 'chat_mention' and state = 'queued'", + ).get(orgId, ceoId) as { id: string }; + await execution.executeRun(orgId, wake.id); + + expect(seenFollowUps).toHaveLength(1); + expect(seenFollowUps[0]).toContain("Messages that arrived while you were working"); + expect(seenFollowUps[0]).toContain("Stop — use Postgres instead."); + expect(seenFollowUps[0]).not.toContain("Use MySQL."); + connection.close(); + }); + + test("a correction that lands while the model answers gets one more turn", async () => { + const h = await founderLine(); + const threadId = h.say("Use MySQL.", null); + h.onWakeTurn((turnIndex) => { + if (turnIndex === 1) h.say("Stop — use Postgres instead.", threadId); + }); + await h.runWake(); + const second = h.wakeInputs[1]; + expect(second.at(-2)).toMatchObject({ role: "assistant", text: "Reply 1." }); + const last = second.at(-1)!; + expect(last.role).toBe("user"); + expect(last.role === "user" && last.text).toContain("Stop — use Postgres instead."); + h.connection.close(); + }); }); diff --git a/tests/unit/provider-tool-turns.test.ts b/tests/unit/provider-tool-turns.test.ts index 2d43dcec..313386cb 100644 --- a/tests/unit/provider-tool-turns.test.ts +++ b/tests/unit/provider-tool-turns.test.ts @@ -281,3 +281,35 @@ describe("responses tool definitions", () => { expect(responsesToolDefs(tools, true)[0]).toMatchObject({ strict: true }); }); }); + +describe("mid-run follow-up text rides on the tool-result turn", () => { + const convo: AgentTurnMessage[] = [ + { role: "user", text: "go" }, + { role: "assistant", text: "", toolCalls: [{ id: "c1", name: "f", arguments: {} }] }, + { role: "tool", results: [{ toolCallId: "c1", toolName: "f", output: "ok" }], followUp: "New: stop" }, + ]; + + test("each wire format keeps the result and appends the text in order", () => { + expect(toResponsesInput(convo).slice(-2)).toEqual([ + { type: "function_call_output", call_id: "c1", output: "ok" }, + { role: "user", content: [{ type: "input_text", text: "New: stop" }] }, + ]); + expect(toChatMessages(null, convo).slice(-2)).toEqual([ + { role: "tool", tool_call_id: "c1", content: "ok" }, + { role: "user", content: "New: stop" }, + ]); + expect(toAnthropicMessages(convo).at(-1)).toEqual({ + role: "user", + content: [{ type: "tool_result", tool_use_id: "c1", content: "ok" }, { type: "text", text: "New: stop" }], + }); + expect(toBedrockMessages(convo).at(-1)).toEqual({ + role: "user", + content: [{ toolResult: { toolUseId: "c1", content: [{ text: "ok" }] } }, { text: "New: stop" }], + }); + const gemini = vertexRequestBody({ model: "gemini-3-pro", instructions: null, messages: convo, tools: [] }) as { contents: Array<{ role: string; parts: unknown[] }> }; + expect(gemini.contents.at(-1)).toEqual({ + role: "user", + parts: [{ functionResponse: { name: "f", response: { output: "ok" } } }, { text: "New: stop" }], + }); + }); +});