getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-17 04:02:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix bash execution interleaving with tool calls: defer message insertion while streaming

Browse source
This commit is contained in:
Mario Zechner 2025-12-09 01:34:31 +01:00
parent 14d99b5f86
commit 2c014c1b95
1 changed files with 67 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

114
packages/coding-agent/src/core/agent-session.ts
View file

@ -107,6 +107,7 @@ export class AgentSession {
// Bash execution state // Bash execution state
private _bashAbortController: AbortController | null = null; private _bashAbortController: AbortController | null = null;
private _pendingBashMessages: BashExecutionMessage[] = [];
constructor(config: AgentSessionConfig) { constructor(config: AgentSessionConfig) {
this.agent = config.agent; this.agent = config.agent;
@ -120,6 +121,34 @@ export class AgentSession {
// Event Subscription // Event Subscription
// ========================================================================= // =========================================================================
/** Internal handler for agent events - shared by subscribe and reconnect */
private _handleAgentEvent = async (event: AgentEvent): Promise<void> => {
// Notify all listeners
for (const l of this._eventListeners) {
l(event);
}
// Handle session persistence
if (event.type === "message_end") {
this.sessionManager.saveMessage(event.message);
// Initialize session after first user+assistant exchange
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
// Check auto-compaction after assistant messages
if (event.message.role === "assistant") {
await this.checkAutoCompaction();
}
}
// Flush pending bash messages after agent turn completes
if (event.type === "agent_end") {
this._flushPendingBashMessages();
}
};
/** /**
* Subscribe to agent events. * Subscribe to agent events.
* Session persistence is handled internally (saves messages on message_end). * Session persistence is handled internally (saves messages on message_end).
@ -130,27 +159,7 @@ export class AgentSession {
// Set up agent subscription if not already done // Set up agent subscription if not already done
if (!this._unsubscribeAgent) { if (!this._unsubscribeAgent) {
this._unsubscribeAgent = this.agent.subscribe(async (event) => { this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
// Notify all listeners
for (const l of this._eventListeners) {
l(event);
}
// Handle session persistence
if (event.type === "message_end") {
this.sessionManager.saveMessage(event.message);
// Initialize session after first user+assistant exchange
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
// Check auto-compaction after assistant messages
if (event.message.role === "assistant") {
await this.checkAutoCompaction();
}
}
});
} }
// Return unsubscribe function for this specific listener // Return unsubscribe function for this specific listener
@ -180,25 +189,7 @@ export class AgentSession {
*/ */
private _reconnectToAgent(): void { private _reconnectToAgent(): void {
if (this._unsubscribeAgent) return; // Already connected if (this._unsubscribeAgent) return; // Already connected
this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
this._unsubscribeAgent = this.agent.subscribe(async (event) => {
for (const l of this._eventListeners) {
l(event);
}
if (event.type === "message_end") {
this.sessionManager.saveMessage(event.message);
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
// Check auto-compaction after assistant messages
if (event.message.role === "assistant") {
await this.checkAutoCompaction();
}
}
});
} }
/** /**
@ -673,15 +664,21 @@ export class AgentSession {
timestamp: Date.now(), timestamp: Date.now(),
}; };
// Add to agent state // If agent is streaming, defer adding to avoid breaking tool_use/tool_result ordering
this.agent.appendMessage(bashMessage); if (this.isStreaming) {
// Queue for later - will be flushed on agent_end
this._pendingBashMessages.push(bashMessage);
} else {
// Add to agent state immediately
this.agent.appendMessage(bashMessage);
// Save to session // Save to session
this.sessionManager.saveMessage(bashMessage); this.sessionManager.saveMessage(bashMessage);
// Initialize session if needed // Initialize session if needed
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) { if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state); this.sessionManager.startSession(this.agent.state);
}
} }
return result; return result;
@ -702,6 +699,29 @@ export class AgentSession {
return this._bashAbortController !== null; return this._bashAbortController !== null;
} }
/**
* Flush pending bash messages to agent state and session.
* Called after agent turn completes to maintain proper message ordering.
*/
private _flushPendingBashMessages(): void {
if (this._pendingBashMessages.length === 0) return;
for (const bashMessage of this._pendingBashMessages) {
// Add to agent state
this.agent.appendMessage(bashMessage);
// Save to session
this.sessionManager.saveMessage(bashMessage);
}
// Initialize session if needed
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
this.sessionManager.startSession(this.agent.state);
}
this._pendingBashMessages = [];
}
// ========================================================================= // =========================================================================
// Session Management // Session Management
// ========================================================================= // =========================================================================