jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Refine agent list attention handling

Browse Source
This commit is contained in:
Mohamed Boudra
2026-03-08 15:39:40 +07:00
parent 03e1915316
commit faa5aaab6f
21 changed files with 1178 additions and 390 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/server/src/server/agent/agent-management-mcp.ts
+13 -41
View File
@@ -153,9 +153,16 @@ function startAgentRun(
agentManager: AgentManager,
agentId: string,
prompt: AgentPromptInput,
logger: Logger
logger: Logger,
options?: { replaceRunning?: boolean }
): void {
const iterator = agentManager.streamAgent(agentId, prompt);
const snapshot = agentManager.getAgent(agentId);
const shouldReplace =
options?.replaceRunning &&
Boolean(snapshot && (snapshot.lifecycle === "running" || snapshot.pendingRun));
const iterator = shouldReplace
? agentManager.replaceAgentRun(agentId, prompt)
: agentManager.streamAgent(agentId, prompt);
void (async () => {
try {
for await (const _ of iterator) {
@@ -565,44 +572,7 @@ export async function createAgentManagementMcpServer(
}
if (snapshot.lifecycle === "running" || snapshot.pendingRun) {
childLogger.debug(
{ agentId },
"Interrupting active run before sending new prompt"
);
try {
const cancelled = await agentManager.cancelAgentRun(agentId);
if (!cancelled) {
childLogger.warn(
{ agentId },
"Agent reported running but no active run was cancelled"
);
}
waitTracker.cancel(agentId, "Agent run interrupted by new prompt");
const maxWaitMs = 5000;
const pollIntervalMs = 50;
const startTime = Date.now();
while (Date.now() - startTime < maxWaitMs) {
const current = agentManager.getAgent(agentId);
if (!current) {
throw new Error(
`Agent ${agentId} not found during cancellation wait`
);
}
if (current.lifecycle !== "running" && !current.pendingRun) {
break;
}
await new Promise((resolve) =>
setTimeout(resolve, pollIntervalMs)
);
}
} catch (error) {
childLogger.error(
{ err: error, agentId },
"Failed to interrupt agent"
);
throw error;
}
waitTracker.cancel(agentId, "Agent run interrupted by new prompt");
}
if (sessionMode) {
@@ -620,7 +590,9 @@ export async function createAgentManagementMcpServer(
);
}
startAgentRun(agentManager, agentId, prompt, childLogger);
startAgentRun(agentManager, agentId, prompt, childLogger, {
replaceRunning: true,
});
if (!background) {
const result = await waitForAgentWithTimeout(agentManager, agentId, {
packages/server/src/server/agent/agent-manager.test.ts
+109
View File
@@ -1029,6 +1029,115 @@ describe("AgentManager", () => {
await consumePromise;
});
test("replaceAgentRun does not emit idle or resolve waiters between interrupted and replacement runs", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-replace-run-"));
const storagePath = join(workdir, "agents");
const storage = new AgentStorage(storagePath, logger);
const allowFirstRunToEnd = deferred<void>();
const allowSecondRunToEnd = deferred<void>();
class ReplaceRunSession extends TestAgentSession {
private streamCount = 0;
override async *stream(): AsyncGenerator<AgentStreamEvent> {
this.streamCount += 1;
if (this.streamCount === 1) {
yield { type: "turn_started", provider: this.provider };
await allowFirstRunToEnd.promise;
yield { type: "turn_completed", provider: this.provider };
return;
}
yield { type: "turn_started", provider: this.provider };
await allowSecondRunToEnd.promise;
yield { type: "turn_completed", provider: this.provider };
}
override async interrupt(): Promise<void> {
allowFirstRunToEnd.resolve();
}
}
class ReplaceRunClient extends TestAgentClient {
override async createSession(config: AgentSessionConfig): Promise<AgentSession> {
return new ReplaceRunSession(config);
}
}
const manager = new AgentManager({
clients: {
codex: new ReplaceRunClient(),
},
registry: storage,
logger,
idFactory: () => "00000000-0000-4000-8000-000000000125",
});
const snapshot = await manager.createAgent({
provider: "codex",
cwd: workdir,
});
const lifecycleUpdates: string[] = [];
const unsubscribe = manager.subscribe(
(event) => {
if (event.type !== "agent_state" || event.agent.id !== snapshot.id) {
return;
}
lifecycleUpdates.push(event.agent.lifecycle);
},
{ agentId: snapshot.id, replayState: false }
);
const firstRun = manager.streamAgent(snapshot.id, "first run");
const firstRunDrain = (async () => {
for await (const _event of firstRun) {
// Drain events so lifecycle updates are applied.
}
})();
await manager.waitForAgentRunStart(snapshot.id);
const waitPromise = manager.waitForAgentEvent(snapshot.id);
const secondRun = manager.replaceAgentRun(snapshot.id, "second run");
const secondRunDrain = (async () => {
for await (const _event of secondRun) {
// Drain replacement run.
}
})();
await manager.waitForAgentRunStart(snapshot.id);
const prematureResolution = await Promise.race([
waitPromise.then(() => "resolved"),
new Promise<"pending">((resolve) => setTimeout(() => resolve("pending"), 50)),
]);
expect(prematureResolution).toBe("pending");
const runningIndexes = lifecycleUpdates.reduce<number[]>((indexes, status, index) => {
if (status === "running") {
indexes.push(index);
}
return indexes;
}, []);
expect(runningIndexes.length).toBeGreaterThanOrEqual(2);
const firstReplacementRunningIndex = runningIndexes[1]!;
expect(
lifecycleUpdates.slice(0, firstReplacementRunningIndex).includes("idle")
).toBe(false);
allowSecondRunToEnd.resolve();
const waited = await waitPromise;
expect(waited.status).toBe("idle");
await firstRunDrain;
await secondRunDrain;
unsubscribe();
});
test("applies live autonomous events while no foreground run is active", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-live-events-"));
const storagePath = join(workdir, "agents");
packages/server/src/server/agent/agent-manager.ts
+62 -6
View File
@@ -153,6 +153,7 @@ type ManagedAgentBase = {
availableModes: AgentMode[];
currentModeId: string | null;
pendingPermissions: Map<string, AgentPermissionRequest>;
pendingReplacement: boolean;
timeline: AgentTimelineItem[];
timelineRows: AgentTimelineRow[];
timelineEpoch: string;
@@ -1038,6 +1039,7 @@ export class AgentManager {
const agent = existingAgent as ActiveManagedAgent;
const iterator = agent.session.stream(prompt, options);
agent.pendingReplacement = false;
agent.lastError = undefined;
let finalized = false;
@@ -1076,7 +1078,13 @@ export class AgentManager {
const mutableAgent = agent as ActiveManagedAgent;
mutableAgent.pendingRun = null;
const terminalError = error ?? mutableAgent.lastError;
mutableAgent.lifecycle = terminalError ? "error" : "idle";
const shouldHoldBusyForReplacement =
mutableAgent.pendingReplacement && !terminalError;
mutableAgent.lifecycle = shouldHoldBusyForReplacement
? "running"
: terminalError
? "error"
: "idle";
mutableAgent.lastError = terminalError;
const persistenceHandle =
mutableAgent.session.describePersistence() ??
@@ -1095,11 +1103,14 @@ export class AgentManager {
lifecycle: mutableAgent.lifecycle,
hasPendingRun: Boolean(mutableAgent.pendingRun),
terminalError,
pendingReplacement: mutableAgent.pendingReplacement,
},
"streamAgent.finalize: applying terminal state"
);
this.emitState(mutableAgent);
this.flushLiveEventBacklog(mutableAgent);
if (!shouldHoldBusyForReplacement) {
this.emitState(mutableAgent);
this.flushLiveEventBacklog(mutableAgent);
}
};
const self = this;
@@ -1162,17 +1173,61 @@ export class AgentManager {
return streamForwarder;
}
replaceAgentRun(
agentId: string,
prompt: AgentPromptInput,
options?: AgentRunOptions
): AsyncGenerator<AgentStreamEvent> {
const snapshot = this.requireAgent(agentId);
if (snapshot.lifecycle !== "running" && !snapshot.pendingRun) {
return this.streamAgent(agentId, prompt, options);
}
const agent = snapshot as ActiveManagedAgent;
agent.pendingReplacement = true;
const self = this;
return (async function* replaceRunForwarder() {
try {
await self.cancelAgentRun(agentId);
const nextRun = self.streamAgent(agentId, prompt, options);
for await (const event of nextRun) {
yield event;
}
} catch (error) {
const latest = self.agents.get(agentId);
if (latest) {
const latestActive = latest as ActiveManagedAgent;
latestActive.pendingReplacement = false;
const hasForegroundRun = Boolean(
(latestActive as { pendingRun: AsyncGenerator<AgentStreamEvent> | null }).pendingRun
);
const lifecycle = (latestActive as { lifecycle: AgentLifecycleStatus }).lifecycle;
if (!hasForegroundRun && lifecycle === "running") {
latestActive.lifecycle = "idle";
self.emitState(latestActive);
self.flushLiveEventBacklog(latestActive);
}
}
throw error;
}
})();
}
async waitForAgentRunStart(agentId: string, options?: WaitForAgentStartOptions): Promise<void> {
const snapshot = this.getAgent(agentId);
if (!snapshot) {
throw new Error(`Agent ${agentId} not found`);
}
if (snapshot.lifecycle === "running") {
if (snapshot.lifecycle === "running" && !snapshot.pendingReplacement) {
return;
}
if (!("pendingRun" in snapshot) || !snapshot.pendingRun) {
if (
(!("pendingRun" in snapshot) || !snapshot.pendingRun) &&
!snapshot.pendingReplacement
) {
throw new Error(`Agent ${agentId} has no pending run`);
}
@@ -1229,7 +1284,7 @@ export class AgentManager {
if (event.agent.id !== agentId) {
return;
}
if (event.agent.lifecycle === "running") {
if (event.agent.lifecycle === "running" && !event.agent.pendingReplacement) {
finishOk();
return;
}
@@ -1575,6 +1630,7 @@ export class AgentManager {
availableModes: [],
currentModeId: null,
pendingPermissions: new Map(),
pendingReplacement: false,
pendingRun: null,
timeline: initialTimeline,
timelineRows: initialTimelineRows,
packages/server/src/server/agent/mcp-server.ts
+13 -43
View File
@@ -225,9 +225,16 @@ function startAgentRun(
agentManager: AgentManager,
agentId: string,
prompt: AgentPromptInput,
logger: Logger
logger: Logger,
options?: { replaceRunning?: boolean }
): void {
const iterator = agentManager.streamAgent(agentId, prompt);
const snapshot = agentManager.getAgent(agentId);
const shouldReplace =
options?.replaceRunning &&
Boolean(snapshot && (snapshot.lifecycle === "running" || snapshot.pendingRun));
const iterator = shouldReplace
? agentManager.replaceAgentRun(agentId, prompt)
: agentManager.streamAgent(agentId, prompt);
void (async () => {
try {
for await (const _ of iterator) {
@@ -764,52 +771,13 @@ export async function createAgentMcpServer(
},
},
async ({ agentId, prompt, sessionMode, background = false }) => {
// Check if agent is running and interrupt if necessary (matches app behavior)
const snapshot = agentManager.getAgent(agentId);
if (!snapshot) {
throw new Error(`Agent ${agentId} not found`);
}
if (snapshot.lifecycle === "running" || snapshot.pendingRun) {
childLogger.debug(
{ agentId },
"Interrupting active run before sending new prompt"
);
try {
const cancelled = await agentManager.cancelAgentRun(agentId);
if (!cancelled) {
childLogger.warn(
{ agentId },
"Agent reported running but no active run was cancelled"
);
}
// Also cancel any pending wait_for_agent calls for this agent
waitTracker.cancel(agentId, "Agent run interrupted by new prompt");
// Wait for the agent to become idle after cancellation
// Poll until the agent is no longer running and has no pending run
// This is necessary because cancelAgentRun only initiates cancellation
// and doesn't wait for the generator to fully terminate
const maxWaitMs = 5000;
const pollIntervalMs = 50;
const startTime = Date.now();
while (Date.now() - startTime < maxWaitMs) {
const current = agentManager.getAgent(agentId);
if (!current) {
throw new Error(`Agent ${agentId} not found during cancellation wait`);
}
if (current.lifecycle !== "running" && !current.pendingRun) {
break;
}
await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
}
} catch (error) {
childLogger.error(
{ err: error, agentId },
"Failed to interrupt agent"
);
throw error;
}
waitTracker.cancel(agentId, "Agent run interrupted by new prompt");
}
if (sessionMode) {
@@ -827,7 +795,9 @@ export async function createAgentMcpServer(
);
}
startAgentRun(agentManager, agentId, prompt, childLogger);
startAgentRun(agentManager, agentId, prompt, childLogger, {
replaceRunning: true,
});
// If not running in background, wait for completion
if (!background) {
packages/server/src/server/daemon-e2e/wait-for-idle.e2e.test.ts
+47
View File
@@ -6,7 +6,9 @@ import {
createDaemonTestContext,
type DaemonTestContext,
} from "../test-utils/index.js";
import { DaemonClient } from "../test-utils/index.js";
import { createMessageCollector, type MessageCollector } from "../test-utils/message-collector.js";
import { getFullAccessConfig } from "./agent-configs.js";
function tmpCwd(): string {
return mkdtempSync(path.join(tmpdir(), "wait-for-idle-e2e-"));
@@ -186,4 +188,49 @@ describe("waitForFinish edge cases", () => {
rmSync(cwd, { recursive: true, force: true });
}
}, 60000);
test("waitForFinish stays blocked when sendMessage transactionally replaces a running turn", async () => {
const cwd = tmpCwd();
const secondary = new DaemonClient({ url: `ws://127.0.0.1:${ctx.daemon.port}/ws` });
const agent = await ctx.client.createAgent({
...getFullAccessConfig("codex"),
cwd,
title: "Transactional Replacement Wait",
});
try {
await secondary.connect();
await secondary.fetchAgents({ subscribe: { subscriptionId: "secondary" } });
await ctx.client.sendMessage(agent.id, "Run: sleep 5");
await ctx.client.waitForAgentUpsert(
agent.id,
(snapshot) => snapshot.status === "running",
5_000
);
const waitPromise = ctx.client.waitForFinish(agent.id, 5_000);
await secondary.sendMessage(agent.id, "Run: sleep 5");
await secondary.waitForAgentUpsert(
agent.id,
(snapshot) => snapshot.status === "running",
5_000
);
const prematureResolution = await Promise.race([
waitPromise.then(() => "resolved"),
new Promise<"pending">((resolve) => setTimeout(() => resolve("pending"), 100)),
]);
expect(prematureResolution).toBe("pending");
const finalState = await waitPromise;
expect(finalState.status).toBe("idle");
} finally {
await secondary.close();
await ctx.client.deleteAgent(agent.id);
rmSync(cwd, { recursive: true, force: true });
}
}, 30000);
});
packages/server/src/server/session.ts
+11 -16
View File
@@ -844,8 +844,15 @@ export class Session {
)
let iterator: AsyncGenerator<AgentStreamEvent>
try {
iterator = this.agentManager.streamAgent(agentId, prompt, runOptions)
this.sessionLogger.trace({ agentId }, 'startAgentStream: streamAgent returned iterator')
const snapshot = this.agentManager.getAgent(agentId)
const shouldReplace = Boolean(snapshot && (snapshot.lifecycle === 'running' || snapshot.pendingRun))
iterator = shouldReplace
? this.agentManager.replaceAgentRun(agentId, prompt, runOptions)
: this.agentManager.streamAgent(agentId, prompt, runOptions)
this.sessionLogger.trace(
{ agentId, shouldReplace },
'startAgentStream: agent iterator returned'
)
} catch (error) {
this.handleAgentRunError(agentId, error, 'Failed to start agent run')
const message =
@@ -2457,19 +2464,6 @@ export class Session {
return
}
try {
await this.interruptAgentIfRunning(agentId)
} catch (error) {
this.handleAgentRunError(
agentId,
error,
'Failed to interrupt running agent before sending prompt'
)
return
}
const prompt = this.buildAgentPrompt(text, images)
try {
this.agentManager.recordUserMessage(agentId, text, {
messageId,
@@ -2482,6 +2476,8 @@ export class Session {
)
}
const prompt = this.buildAgentPrompt(text, images)
this.startAgentStream(agentId, prompt, runOptions)
}
@@ -6100,7 +6096,6 @@ export class Session {
await this.unarchiveAgentState(agentId)
await this.ensureAgentLoaded(agentId)
await this.interruptAgentIfRunning(agentId)
try {
this.agentManager.recordUserMessage(agentId, msg.text, {