Skip to content

Commit 19cf623

Browse files
committed
feat(agents): enhance execution control with stop functionality and improved database handling
- Refactor database path handling in agents.rs to initialize paths earlier and make them available to all async tasks - Add immediate session ID updates to the database when session IDs are extracted - Implement stop/cancel functionality for running agent sessions in both AgentRunOutputViewer and AgentRunView - Improve loading states and event listener management in AgentRunOutputViewer - Add stop button UI controls with proper styling and positioning - Enhance error handling and logging throughout the agent execution flow
1 parent 85dce56 commit 19cf623

3 files changed

Lines changed: 259 additions & 62 deletions

File tree

src-tauri/src/commands/agents.rs

Lines changed: 53 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -814,6 +814,14 @@ async fn spawn_agent_sidecar(
814814
// We'll extract the session ID from Claude's init message
815815
let session_id_holder: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
816816

817+
// Create variables we need for the spawned task
818+
let app_dir = app
819+
.path()
820+
.app_data_dir()
821+
.expect("Failed to get app data dir");
822+
let db_path = app_dir.join("agents.db");
823+
let db_path_for_stream = db_path.clone(); // Clone for the streaming task
824+
817825
// Spawn task to read events from sidecar
818826
let app_handle = app.clone();
819827
let session_id_holder_clone = session_id_holder.clone();
@@ -865,6 +873,23 @@ async fn spawn_agent_sidecar(
865873
if current_session_id.is_none() {
866874
*current_session_id = Some(sid.to_string());
867875
info!("🔑 Extracted session ID: {}", sid);
876+
877+
// Update database immediately with session ID
878+
if let Ok(conn) = Connection::open(&db_path_for_stream) {
879+
match conn.execute(
880+
"UPDATE agent_runs SET session_id = ?1 WHERE id = ?2",
881+
params![sid, run_id],
882+
) {
883+
Ok(rows) => {
884+
if rows > 0 {
885+
info!("✅ Updated agent run {} with session ID immediately", run_id);
886+
}
887+
}
888+
Err(e) => {
889+
error!("❌ Failed to update session ID immediately: {}", e);
890+
}
891+
}
892+
}
868893
}
869894
}
870895
}
@@ -905,13 +930,6 @@ async fn spawn_agent_sidecar(
905930
info!("📖 Finished reading Claude sidecar events. Total lines: {}", line_count);
906931
});
907932

908-
// Create variables we need for the spawned task
909-
let app_dir = app
910-
.path()
911-
.app_data_dir()
912-
.expect("Failed to get app data dir");
913-
let db_path = app_dir.join("agents.db");
914-
915933
// Monitor process status and wait for completion
916934
tokio::spawn(async move {
917935
info!("🕐 Starting sidecar process monitoring...");
@@ -1040,6 +1058,13 @@ async fn spawn_agent_system(
10401058
let stdout_reader = TokioBufReader::new(stdout);
10411059
let stderr_reader = TokioBufReader::new(stderr);
10421060

1061+
// Create variables we need for the spawned tasks
1062+
let app_dir = app
1063+
.path()
1064+
.app_data_dir()
1065+
.expect("Failed to get app data dir");
1066+
let db_path = app_dir.join("agents.db");
1067+
10431068
// Shared state for collecting session ID and live output
10441069
let session_id = std::sync::Arc::new(Mutex::new(String::new()));
10451070
let live_output = std::sync::Arc::new(Mutex::new(String::new()));
@@ -1052,6 +1077,7 @@ async fn spawn_agent_system(
10521077
let registry_clone = registry.0.clone();
10531078
let first_output = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
10541079
let first_output_clone = first_output.clone();
1080+
let db_path_for_stdout = db_path.clone(); // Clone the db_path for the stdout task
10551081

10561082
let stdout_task = tokio::spawn(async move {
10571083
info!("📖 Starting to read Claude stdout...");
@@ -1095,6 +1121,23 @@ async fn spawn_agent_system(
10951121
if current_session_id.is_empty() {
10961122
*current_session_id = sid.to_string();
10971123
info!("🔑 Extracted session ID: {}", sid);
1124+
1125+
// Update database immediately with session ID
1126+
if let Ok(conn) = Connection::open(&db_path_for_stdout) {
1127+
match conn.execute(
1128+
"UPDATE agent_runs SET session_id = ?1 WHERE id = ?2",
1129+
params![sid, run_id],
1130+
) {
1131+
Ok(rows) => {
1132+
if rows > 0 {
1133+
info!("✅ Updated agent run {} with session ID immediately", run_id);
1134+
}
1135+
}
1136+
Err(e) => {
1137+
error!("❌ Failed to update session ID immediately: {}", e);
1138+
}
1139+
}
1140+
}
10981141
}
10991142
}
11001143
}
@@ -1164,12 +1207,7 @@ async fn spawn_agent_system(
11641207
.map_err(|e| format!("Failed to register process: {}", e))?;
11651208
info!("📋 Registered process in registry");
11661209

1167-
// Create variables we need for the spawned task
1168-
let app_dir = app
1169-
.path()
1170-
.app_data_dir()
1171-
.expect("Failed to get app data dir");
1172-
let db_path = app_dir.join("agents.db");
1210+
let db_path_for_monitor = db_path.clone(); // Clone for the monitor task
11731211

11741212
// Monitor process status and wait for completion
11751213
tokio::spawn(async move {
@@ -1221,7 +1259,7 @@ async fn spawn_agent_system(
12211259
}
12221260

12231261
// Update database
1224-
if let Ok(conn) = Connection::open(&db_path) {
1262+
if let Ok(conn) = Connection::open(&db_path_for_monitor) {
12251263
let _ = conn.execute(
12261264
"UPDATE agent_runs SET status = 'failed', completed_at = CURRENT_TIMESTAMP WHERE id = ?1",
12271265
params![run_id],
@@ -1255,7 +1293,7 @@ async fn spawn_agent_system(
12551293
info!("✅ Claude process execution monitoring complete");
12561294

12571295
// Update the run record with session ID and mark as completed - open a new connection
1258-
if let Ok(conn) = Connection::open(&db_path) {
1296+
if let Ok(conn) = Connection::open(&db_path_for_monitor) {
12591297
info!("🔄 Updating database with extracted session ID: {}", extracted_session_id);
12601298
match conn.execute(
12611299
"UPDATE agent_runs SET session_id = ?1, status = 'completed', completed_at = CURRENT_TIMESTAMP WHERE id = ?2",

src/components/AgentRunOutputViewer.tsx

Lines changed: 109 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ import {
1212
Clock,
1313
Hash,
1414
DollarSign,
15-
ExternalLink
15+
ExternalLink,
16+
StopCircle
1617
} from 'lucide-react';
1718
import { Button } from '@/components/ui/button';
1819
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
@@ -64,13 +65,17 @@ export function AgentRunOutputViewer({
6465
}: AgentRunOutputViewerProps) {
6566
const [messages, setMessages] = useState<ClaudeStreamMessage[]>([]);
6667
const [rawJsonlOutput, setRawJsonlOutput] = useState<string[]>([]);
67-
const [loading, setLoading] = useState(false);
68+
const [loading, setLoading] = useState(true);
6869
const [isFullscreen, setIsFullscreen] = useState(false);
6970
const [refreshing, setRefreshing] = useState(false);
7071
const [toast, setToast] = useState<{ message: string; type: "success" | "error" } | null>(null);
7172
const [copyPopoverOpen, setCopyPopoverOpen] = useState(false);
7273
const [hasUserScrolled, setHasUserScrolled] = useState(false);
7374

75+
// Track whether we're in the initial load phase
76+
const isInitialLoadRef = useRef(true);
77+
const hasSetupListenersRef = useRef(false);
78+
7479
const scrollAreaRef = useRef<HTMLDivElement>(null);
7580
const outputEndRef = useRef<HTMLDivElement>(null);
7681
const fullscreenScrollRef = useRef<HTMLDivElement>(null);
@@ -98,10 +103,12 @@ export function AgentRunOutputViewer({
98103
}
99104
};
100105

101-
// Clean up listeners on unmount
106+
// Cleanup on unmount
102107
useEffect(() => {
103108
return () => {
104109
unlistenRefs.current.forEach(unlisten => unlisten());
110+
unlistenRefs.current = [];
111+
hasSetupListenersRef.current = false;
105112
};
106113
}, []);
107114

@@ -235,35 +242,52 @@ export function AgentRunOutputViewer({
235242
}
236243
};
237244

245+
// Set up live event listeners for running sessions
238246
const setupLiveEventListeners = async () => {
239-
if (!run.id) return;
247+
if (!run.id || hasSetupListenersRef.current) return;
240248

241249
try {
242250
// Clean up existing listeners
243251
unlistenRefs.current.forEach(unlisten => unlisten());
244252
unlistenRefs.current = [];
245253

254+
// Mark that we've set up listeners
255+
hasSetupListenersRef.current = true;
256+
257+
// After setup, we're no longer in initial load
258+
// Small delay to ensure any pending messages are processed
259+
setTimeout(() => {
260+
isInitialLoadRef.current = false;
261+
}, 100);
262+
246263
// Set up live event listeners with run ID isolation
247264
const outputUnlisten = await listen<string>(`agent-output:${run.id}`, (event) => {
248265
try {
266+
// Skip messages during initial load phase
267+
if (isInitialLoadRef.current) {
268+
console.log('[AgentRunOutputViewer] Skipping message during initial load');
269+
return;
270+
}
271+
249272
// Store raw JSONL
250273
setRawJsonlOutput(prev => [...prev, event.payload]);
251274

252275
// Parse and display
253276
const message = JSON.parse(event.payload) as ClaudeStreamMessage;
254277
setMessages(prev => [...prev, message]);
255278
} catch (err) {
256-
console.error("Failed to parse message:", err, event.payload);
279+
console.error("[AgentRunOutputViewer] Failed to parse message:", err, event.payload);
257280
}
258281
});
259282

260283
const errorUnlisten = await listen<string>(`agent-error:${run.id}`, (event) => {
261-
console.error("Agent error:", event.payload);
284+
console.error("[AgentRunOutputViewer] Agent error:", event.payload);
262285
setToast({ message: event.payload, type: 'error' });
263286
});
264287

265288
const completeUnlisten = await listen<boolean>(`agent-complete:${run.id}`, () => {
266289
setToast({ message: 'Agent execution completed', type: 'success' });
290+
// Don't set status here as the parent component should handle it
267291
});
268292

269293
const cancelUnlisten = await listen<boolean>(`agent-cancelled:${run.id}`, () => {
@@ -272,7 +296,7 @@ export function AgentRunOutputViewer({
272296

273297
unlistenRefs.current = [outputUnlisten, errorUnlisten, completeUnlisten, cancelUnlisten];
274298
} catch (error) {
275-
console.error('Failed to set up live event listeners:', error);
299+
console.error('[AgentRunOutputViewer] Failed to set up live event listeners:', error);
276300
}
277301
};
278302

@@ -341,12 +365,63 @@ export function AgentRunOutputViewer({
341365
setToast({ message: 'Output copied as Markdown', type: 'success' });
342366
};
343367

344-
const refreshOutput = async () => {
368+
const handleRefresh = async () => {
345369
setRefreshing(true);
346-
await loadOutput(true); // Skip cache
370+
await loadOutput();
347371
setRefreshing(false);
348372
};
349373

374+
const handleStop = async () => {
375+
if (!run.id) {
376+
console.error('[AgentRunOutputViewer] No run ID available to stop');
377+
return;
378+
}
379+
380+
try {
381+
// Call the API to kill the agent session
382+
const success = await api.killAgentSession(run.id);
383+
384+
if (success) {
385+
console.log(`[AgentRunOutputViewer] Successfully stopped agent session ${run.id}`);
386+
setToast({ message: 'Agent execution stopped', type: 'success' });
387+
388+
// Clean up listeners
389+
unlistenRefs.current.forEach(unlisten => unlisten());
390+
unlistenRefs.current = [];
391+
hasSetupListenersRef.current = false;
392+
393+
// Add a message indicating execution was stopped
394+
const stopMessage: ClaudeStreamMessage = {
395+
type: "result",
396+
subtype: "error",
397+
is_error: true,
398+
result: "Execution stopped by user",
399+
duration_ms: 0,
400+
usage: {
401+
input_tokens: 0,
402+
output_tokens: 0
403+
}
404+
};
405+
setMessages(prev => [...prev, stopMessage]);
406+
407+
// Update the run status locally
408+
// Optionally refresh the parent component
409+
setTimeout(() => {
410+
window.location.reload(); // Simple refresh to update the status
411+
}, 1000);
412+
} else {
413+
console.warn(`[AgentRunOutputViewer] Failed to stop agent session ${run.id} - it may have already finished`);
414+
setToast({ message: 'Failed to stop agent - it may have already finished', type: 'error' });
415+
}
416+
} catch (err) {
417+
console.error('[AgentRunOutputViewer] Failed to stop agent:', err);
418+
setToast({
419+
message: `Failed to stop execution: ${err instanceof Error ? err.message : 'Unknown error'}`,
420+
type: 'error'
421+
});
422+
}
423+
};
424+
350425
const handleScroll = (e: React.UIEvent<HTMLDivElement>) => {
351426
const target = e.currentTarget;
352427
const { scrollTop, scrollHeight, clientHeight } = target;
@@ -562,13 +637,25 @@ export function AgentRunOutputViewer({
562637
<Button
563638
variant="ghost"
564639
size="sm"
565-
onClick={refreshOutput}
640+
onClick={handleRefresh}
566641
disabled={refreshing}
567642
title="Refresh output"
568643
className="h-8 px-2"
569644
>
570645
<RotateCcw className={`h-4 w-4 ${refreshing ? 'animate-spin' : ''}`} />
571646
</Button>
647+
{run.status === 'running' && (
648+
<Button
649+
variant="ghost"
650+
size="sm"
651+
onClick={handleStop}
652+
disabled={refreshing}
653+
title="Stop execution"
654+
className="h-8 px-2 text-destructive hover:text-destructive"
655+
>
656+
<StopCircle className="h-4 w-4" />
657+
</Button>
658+
)}
572659
<Button
573660
variant="ghost"
574661
size="sm"
@@ -667,11 +754,22 @@ export function AgentRunOutputViewer({
667754
<Button
668755
variant="outline"
669756
size="sm"
670-
onClick={refreshOutput}
757+
onClick={handleRefresh}
671758
disabled={refreshing}
672759
>
673760
<RotateCcw className={`h-4 w-4 ${refreshing ? 'animate-spin' : ''}`} />
674761
</Button>
762+
{run.status === 'running' && (
763+
<Button
764+
variant="outline"
765+
size="sm"
766+
onClick={handleStop}
767+
disabled={refreshing}
768+
>
769+
<StopCircle className="h-4 w-4 mr-2" />
770+
Stop
771+
</Button>
772+
)}
675773
<Button
676774
variant="outline"
677775
size="sm"

0 commit comments

Comments
 (0)