Skip to content

Commit 9ff6cd0

Browse files
milispclaude
andcommitted
feat: add streaming control with stop functionality
- Implement interrupt functionality for codex sessions - Add pause_session command alongside existing stop_session - Update ChatInput with stop button during streaming - Add stop streaming handler in ChatInterface - Improve session state management for loading states 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 557a25c commit 9ff6cd0

File tree

7 files changed

+86
-18
lines changed

7 files changed

+86
-18
lines changed

src-tauri/src/codex_client.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,6 @@ impl CodexClient {
273273
self.send_submission(submission).await
274274
}
275275

276-
#[allow(dead_code)]
277276
pub async fn interrupt(&self) -> Result<()> {
278277
let submission = Submission {
279278
id: Uuid::new_v4().to_string(),
@@ -338,7 +337,8 @@ impl CodexClient {
338337
log::debug!("Session {} closed", self.session_id);
339338
Ok(())
340339
}
341-
340+
341+
#[allow(dead_code)]
342342
pub async fn shutdown(&mut self) -> Result<()> {
343343
self.close_session().await
344344
}

src-tauri/src/commands.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,11 @@ pub async fn stop_session(state: State<'_, CodexState>, session_id: String) -> R
4646
codex::stop_session(state, session_id).await
4747
}
4848

49+
#[tauri::command]
50+
pub async fn pause_session(state: State<'_, CodexState>, session_id: String) -> Result<(), String> {
51+
codex::pause_session(state, session_id).await
52+
}
53+
4954
#[tauri::command]
5055
pub async fn close_session(state: State<'_, CodexState>, session_id: String) -> Result<(), String> {
5156
codex::close_session(state, session_id).await

src-tauri/src/lib.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ mod utils;
99

1010
use commands::{
1111
approve_execution, check_codex_version, close_session, delete_session_file,
12-
get_latest_session_id, get_running_sessions, load_sessions_from_disk, send_message,
12+
get_latest_session_id, get_running_sessions, load_sessions_from_disk, pause_session, send_message,
1313
start_codex_session, stop_session,
1414
};
1515
use config::{
@@ -53,6 +53,7 @@ pub fn run() {
5353
send_message,
5454
approve_execution,
5555
stop_session,
56+
pause_session,
5657
close_session,
5758
get_running_sessions,
5859
load_sessions_from_disk,

src-tauri/src/services/codex.rs

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,18 +74,38 @@ pub async fn approve_execution(
7474
}
7575

7676
pub async fn stop_session(state: State<'_, CodexState>, session_id: String) -> Result<(), String> {
77-
let mut sessions = state.sessions.lock().await;
77+
let sessions = state.sessions.lock().await;
7878
let stored_sessions: Vec<String> = sessions.keys().cloned().collect();
7979

80-
log::debug!("Attempting to stop session: {}", session_id);
80+
log::debug!("Attempting to interrupt session: {}", session_id);
8181
log::debug!("Currently stored sessions: {:?}", stored_sessions);
8282

83-
if let Some(mut client) = sessions.remove(&session_id) {
84-
log::debug!("Found and removing session: {}", session_id);
83+
if let Some(client) = sessions.get(&session_id) {
84+
log::debug!("Found session, sending interrupt: {}", session_id);
85+
client
86+
.interrupt()
87+
.await
88+
.map_err(|e| format!("Failed to interrupt session: {}", e))?;
89+
Ok(())
90+
} else {
91+
log::debug!("Session not found: {}", session_id);
92+
Err("Session not found".to_string())
93+
}
94+
}
95+
96+
pub async fn pause_session(state: State<'_, CodexState>, session_id: String) -> Result<(), String> {
97+
let sessions = state.sessions.lock().await;
98+
let stored_sessions: Vec<String> = sessions.keys().cloned().collect();
99+
100+
log::debug!("Attempting to pause session: {}", session_id);
101+
log::debug!("Currently stored sessions: {:?}", stored_sessions);
102+
103+
if let Some(client) = sessions.get(&session_id) {
104+
log::debug!("Found session, sending interrupt (pause): {}", session_id);
85105
client
86-
.shutdown()
106+
.interrupt()
87107
.await
88-
.map_err(|e| format!("Failed to shutdown session: {}", e))?;
108+
.map_err(|e| format!("Failed to pause session: {}", e))?;
89109
Ok(())
90110
} else {
91111
log::debug!("Session not found: {}", session_id);

src/components/chat/ChatInput.tsx

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import React, { useState, useEffect } from 'react';
22
import { Button } from '../ui/button';
33
import { Textarea } from '../ui/textarea';
44
import { Badge } from '../ui/badge';
5-
import { Send, AtSign, X, ChevronUp, Cpu } from 'lucide-react';
5+
import { Send, AtSign, X, ChevronUp, Cpu, Square } from 'lucide-react';
66
import {
77
Tooltip,
88
TooltipContent,
@@ -23,6 +23,7 @@ interface ChatInputProps {
2323
inputValue: string;
2424
onInputChange: (value: string) => void;
2525
onSendMessage: (message: string) => void;
26+
onStopStreaming?: () => void;
2627
disabled?: boolean;
2728
isLoading?: boolean;
2829
placeholderOverride?: string;
@@ -32,6 +33,7 @@ export const ChatInput: React.FC<ChatInputProps> = ({
3233
inputValue,
3334
onInputChange,
3435
onSendMessage,
36+
onStopStreaming,
3537
disabled = false,
3638
isLoading = false,
3739
placeholderOverride,
@@ -144,6 +146,12 @@ export const ChatInput: React.FC<ChatInputProps> = ({
144146
clearFileReferences();
145147
};
146148

149+
const handleStopStreaming = () => {
150+
if (onStopStreaming) {
151+
onStopStreaming();
152+
}
153+
};
154+
147155
const handleKeyPress = (e: React.KeyboardEvent) => {
148156
if (e.key === 'Enter' && !e.shiftKey) {
149157
e.preventDefault();
@@ -201,14 +209,25 @@ export const ChatInput: React.FC<ChatInputProps> = ({
201209
className="flex-1 min-h-[40px] max-h-[120px]"
202210
disabled={false}
203211
/>
204-
<Button
205-
onClick={handleSendMessage}
206-
disabled={!inputValue.trim() || isLoading || disabled}
207-
size="sm"
208-
className="self-end"
209-
>
210-
<Send className="w-4 h-4" />
211-
</Button>
212+
{isLoading ? (
213+
<Button
214+
onClick={handleStopStreaming}
215+
size="sm"
216+
className="self-end bg-red-500 hover:bg-red-600 text-white"
217+
variant="default"
218+
>
219+
<Square className="w-4 h-4" />
220+
</Button>
221+
) : (
222+
<Button
223+
onClick={handleSendMessage}
224+
disabled={!inputValue.trim() || disabled}
225+
size="sm"
226+
className="self-end"
227+
>
228+
<Send className="w-4 h-4" />
229+
</Button>
230+
)}
212231
</div>
213232

214233
{/* Model Selection Bar */}

src/components/chat/ChatInterface.tsx

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,27 @@ export const ChatInterface: React.FC<ChatInterfaceProps> = ({
314314
}
315315
};
316316

317+
const handleStopStreaming = async () => {
318+
try {
319+
// Extract raw session ID for backend communication
320+
const rawSessionId = sessionId.startsWith('codex-event-')
321+
? sessionId.replace('codex-event-', '')
322+
: sessionId;
323+
324+
await invoke("pause_session", {
325+
sessionId: rawSessionId,
326+
});
327+
328+
// Immediately set loading to false after successful pause
329+
setSessionLoading(sessionId, false);
330+
331+
} catch (error) {
332+
console.error("Failed to pause streaming:", error);
333+
// On error, also set loading to false
334+
setSessionLoading(sessionId, false);
335+
}
336+
};
337+
317338
return (
318339
<div className="flex h-full min-h-0">
319340
{/* Session Manager - conditionally visible */}
@@ -349,6 +370,7 @@ export const ChatInterface: React.FC<ChatInterfaceProps> = ({
349370
inputValue={inputValue}
350371
onInputChange={setInputValue}
351372
onSendMessage={handleSendMessage}
373+
onStopStreaming={handleStopStreaming}
352374
disabled={!!selectedConversation}
353375
isLoading={isLoading}
354376
/>

src/hooks/useCodexEvents.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ export const useCodexEvents = ({
8888
break;
8989

9090
case 'task_complete':
91+
console.log('🔄 Task complete event received, setting loading to false');
9192
setSessionLoading(sessionId, false);
9293
// Finalize any ongoing stream
9394
if (currentStreamingMessageId.current) {

0 commit comments

Comments
 (0)