Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -129,8 +129,6 @@ public enum CLILineToBlocksMapper {
switch part {
case .text(let t):
guard !t.isEmpty else { continue }
let trimmed = t.trimmingCharacters(in: .whitespacesAndNewlines)
if CLIMetaEnvelope.isNoResponseRequested(trimmed) { continue }
blocks.append(.text(t))
case .toolUse(let id, let name, let input):
blocks.append(.toolCall(ToolCall(id: id, name: name, input: input)))
Expand Down
57 changes: 26 additions & 31 deletions RxCode/App/AppState+CrossProject.swift
Original file line number Diff line number Diff line change
Expand Up @@ -105,9 +105,23 @@
}

// After sendPrompt returns, window.currentSessionId is the (possibly
// pending-) key the stream is bound to. The CLI may rename it to its
// own sid mid-stream; we surface whichever id the completion lands on.
let postSendThreadId = window.currentSessionId ?? resolvedThreadId ?? ""
// pending-) key the stream is bound to. Resolve the real CLI session
// id before returning so the caller's agent never sees `pending-…`
// (which it can't use to follow up via `get_thread_messages` etc.).
let postSendKey = window.currentSessionId ?? resolvedThreadId ?? ""
let resolvedThreadIdForReturn: String
if postSendKey.hasPrefix("pending-") {
// Cap the rename wait at the request's timeout so we still honor
// the caller's deadline; 60s upper bound matches typical first-token
// latency under healthy conditions.
let renameTimeout = min(max(timeoutSeconds, 1), 60)
resolvedThreadIdForReturn = await awaitSessionRename(
pendingKey: postSendKey,
timeout: renameTimeout
) ?? postSendKey
} else {
resolvedThreadIdForReturn = postSendKey
}

if !waitForResponse {
// Don't leak the result in the dictionary — the caller is
Expand All @@ -116,7 +130,7 @@
_ = await self?.awaitStreamCompletion(streamId: streamId, timeout: timeoutSeconds)
}
return CrossProjectSendResult(
threadId: postSendThreadId,
threadId: resolvedThreadIdForReturn,
projectId: resolvedProject.id,
done: false,
assistantText: "",
Expand All @@ -138,7 +152,7 @@
// the caller can decide whether to poll back via get_thread_messages.
let partial = lastAssistantResponseText(in: stateForSession(window.currentSessionId ?? "").messages)
return CrossProjectSendResult(
threadId: window.currentSessionId ?? postSendThreadId,
threadId: resolvedThreadIdForReturn,
projectId: resolvedProject.id,
done: false,
assistantText: partial,
Expand All @@ -147,28 +161,6 @@
}
}

/// Drop "No response requested." text blocks from the assistant message
/// at `idx`. If the message has no blocks left after the strip, remove
/// it entirely. Called at turn-finalization sites — the marker is the
/// model's response when a turn arrives without a user prompt
/// (ScheduleWakeup, hook re-entry) and reads as noise in the chat UI.
/// Strip CLI no-op meta text ("no response requested") from a message.
///
/// `removeIfEmpty` controls whether a message left with no blocks is also
/// deleted. The normal stream path passes `true` to discard pure no-op
/// envelopes; the cancel path passes `false` so pausing a turn never makes
/// the partial assistant bubble disappear.
static func stripNoOpText(at idx: Int, in messages: inout [ChatMessage], removeIfEmpty: Bool = true) {
guard messages.indices.contains(idx) else { return }
messages[idx].blocks.removeAll { block in
guard let text = block.text else { return false }
return CLIMetaEnvelope.isNoResponseRequested(text.trimmingCharacters(in: .whitespacesAndNewlines))
}
if removeIfEmpty, messages[idx].blocks.isEmpty {
messages.remove(at: idx)
}
}

/// Wrap a branch briefing into a system-prompt section the agent can use as
/// background context. The briefing is auto-generated from earlier threads,
/// so it is framed as advisory rather than authoritative.
Expand Down Expand Up @@ -395,11 +387,15 @@

var eventCount = 0
var lastEventTime = Date()
logger.info("[Stream:UI] entering for-await session=\(sessionKey, privacy: .public) stream=\(streamId) cwd=\(cwd, privacy: .public)")

do {
for await event in stream {
eventCount += 1
let gap = Date().timeIntervalSince(lastEventTime)
if eventCount == 1 {
logger.info("[Stream:UI] first event arrived session=\(sessionKey, privacy: .public) stream=\(streamId) after=\(String(format: "%.2f", gap))s")
}
lastEventTime = Date()
updateState(sessionKey) { $0.lastStreamEventDate = lastEventTime }

Expand All @@ -418,7 +414,7 @@
if let state = sessionStates.removeValue(forKey: sessionKey) {
sessionStates[resultEvent.sessionId] = state
}
sessionIdRedirect[sessionKey] = resultEvent.sessionId
applySessionIdRedirect(from: sessionKey, to: resultEvent.sessionId)
sessionKey = resultEvent.sessionId
}
let msgs = stateForSession(sessionKey).messages
Expand Down Expand Up @@ -452,7 +448,7 @@
sessionStates[sid] = state
}
renameDraftState(from: previousSessionKey, to: sid, in: window)
sessionIdRedirect[previousSessionKey] = sid
applySessionIdRedirect(from: previousSessionKey, to: sid)
sessionKey = sid
startFlushTimer(for: sid)

Expand Down Expand Up @@ -674,7 +670,6 @@
}) {
state.messages[idx].isStreaming = false
state.messages[idx].finalizeToolCalls()
Self.stripNoOpText(at: idx, in: &state.messages)
}
state.messages.append(ChatMessage(role: .assistant, isStreaming: true))
state.needsNewMessage = false
Expand Down Expand Up @@ -716,7 +711,7 @@
sessionStates[resultEvent.sessionId] = state
}
renameDraftState(from: previousSessionKey, to: resultEvent.sessionId, in: window)
sessionIdRedirect[previousSessionKey] = resultEvent.sessionId
applySessionIdRedirect(from: previousSessionKey, to: resultEvent.sessionId)
sessionKey = resultEvent.sessionId
if wasForeground {
window.currentSessionId = resultEvent.sessionId
Expand Down Expand Up @@ -960,4 +955,4 @@
}
}

}

Check warning on line 958 in RxCode/App/AppState+CrossProject.swift

View workflow job for this annotation

GitHub Actions / swiftlint

File should contain 600 lines or less excluding comments and whitespaces: currently contains 775 (file_length)
139 changes: 125 additions & 14 deletions RxCode/App/AppState+Messaging.swift
Original file line number Diff line number Diff line change
Expand Up @@ -456,7 +456,6 @@ extension AppState {
if let start = state.streamingStartDate {
state.messages[idx].duration = Date().timeIntervalSince(start)
}
Self.stripNoOpText(at: idx, in: &state.messages)
}
state.streamingStartDate = nil
}
Expand All @@ -468,42 +467,154 @@ extension AppState {

// MARK: - Stream Completion (cross-project MCP)

/// Record that the stream `streamId` finished. Stored in
/// `pendingStreamCompletions` for any `awaitStreamCompletion(...)` caller
/// (currently `ide__send_to_thread`) to pick up. Latest call wins, except
/// we don't overwrite a success with an error from the fallback path.
/// Record that the stream `streamId` finished. If a caller is already
/// waiting via `awaitStreamCompletion(...)`, the result is handed to it
/// directly so it can return immediately; otherwise it's parked in
/// `pendingStreamCompletions` until someone picks it up.
func recordStreamCompletion(
streamId: UUID,
sessionId: String,
assistantText: String,
error: String?
) {
pendingStreamCompletions[streamId] = StreamCompletion(
let completion = StreamCompletion(
sessionId: sessionId,
assistantText: assistantText,
error: error
)
if let waiter = streamCompletionWaiters.removeValue(forKey: streamId) {
waiter.resume(with: completion)
return
}
pendingStreamCompletions[streamId] = completion
Comment on lines +485 to +489
}

/// Wait up to `timeout` seconds for the stream identified by `streamId`
/// to record a completion. Polls every 100ms — MainActor serialization
/// means the recorder fires between sleeps. Returns the completion if
/// one arrived in time, otherwise `nil`.
/// to record a completion. Event-driven: `recordStreamCompletion` resumes
/// the continuation as soon as the result lands, with a parallel
/// `Task.sleep(timeout)` to surface `nil` if the deadline passes first.
/// Returns the completion if one arrived in time, otherwise `nil`.
func awaitStreamCompletion(streamId: UUID, timeout: TimeInterval) async -> StreamCompletion? {
let deadline = Date().addingTimeInterval(timeout)
while Date() < deadline {
// Fast path — completion already landed before we registered.
if let completion = pendingStreamCompletions.removeValue(forKey: streamId) {
return completion
}

let timeoutTask = Task { [weak self] in
let nanos = UInt64(max(0, timeout) * 1_000_000_000)
try? await Task.sleep(nanoseconds: nanos)
guard let self else { return }
await MainActor.run {
if let waiter = self.streamCompletionWaiters.removeValue(forKey: streamId) {
waiter.resume(with: nil)
}
}
}

let result: StreamCompletion? = await withCheckedContinuation { cont in
let waiter = StreamCompletionWaiter(continuation: cont)
// Re-check between the fast-path read and continuation install — the
// recorder may have fired in between if any MainActor work yielded.
if let completion = pendingStreamCompletions.removeValue(forKey: streamId) {
return completion
waiter.resume(with: completion)
} else {
streamCompletionWaiters[streamId] = waiter
}
try? await Task.sleep(nanoseconds: 100_000_000)
}
return pendingStreamCompletions.removeValue(forKey: streamId)

timeoutTask.cancel()
return result
}

/// Discard a recorded completion. Called by long-running `wait_for_response=false`
/// MCP sends so the dictionary doesn't grow unbounded with abandoned results.
func discardStreamCompletion(streamId: UUID) {
pendingStreamCompletions.removeValue(forKey: streamId)
if let waiter = streamCompletionWaiters.removeValue(forKey: streamId) {
Comment on lines 529 to +533
waiter.resume(with: nil)
}
}

// MARK: - Session-id rename handoff (cross-project MCP)

/// Record that `pendingKey` was renamed to `realSessionId` (the CLI's own
/// `session_id`). Mirror writes to `sessionIdRedirect` already; this also
/// resumes any `awaitSessionRename` caller so the cross-project send can
/// return the real thread id instead of `pending-…`.
func applySessionIdRedirect(from pendingKey: String, to realSessionId: String) {
sessionIdRedirect[pendingKey] = realSessionId
if let waiter = sessionIdRenameWaiters.removeValue(forKey: pendingKey) {
waiter.resume(with: realSessionId)
}
}

/// Wait up to `timeout` seconds for `pendingKey` to be renamed to the
/// CLI's real `session_id`. Fast-paths the answer if the rename already
/// landed. Returns `nil` on timeout.
func awaitSessionRename(pendingKey: String, timeout: TimeInterval) async -> String? {
if let real = sessionIdRedirect[pendingKey] {
return real
}

let timeoutTask = Task { [weak self] in
let nanos = UInt64(max(0, timeout) * 1_000_000_000)
try? await Task.sleep(nanoseconds: nanos)
guard let self else { return }
await MainActor.run {
if let waiter = self.sessionIdRenameWaiters.removeValue(forKey: pendingKey) {
waiter.resume(with: nil)
}
}
}

let result: String? = await withCheckedContinuation { cont in
let waiter = SessionRenameWaiter(continuation: cont)
if let real = sessionIdRedirect[pendingKey] {
waiter.resume(with: real)
} else {
sessionIdRenameWaiters[pendingKey] = waiter
}
}

timeoutTask.cancel()
return result
}

}

/// Resumes a single waiting `awaitSessionRename` caller. Same single-shot
/// pattern as `StreamCompletionWaiter` — guards against double-resume in the
/// race between the rename notification and the timeout task.
@MainActor
final class SessionRenameWaiter {
private var continuation: CheckedContinuation<String?, Never>?

init(continuation: CheckedContinuation<String?, Never>) {
self.continuation = continuation
}

func resume(with value: String?) {
guard let cont = continuation else { return }
continuation = nil
cont.resume(returning: value)
}
}

/// Resumes a single waiting `awaitStreamCompletion` caller. The class wrapper
/// guards against double-resume in the race between `recordStreamCompletion`
/// and the timeout task: whichever fires first wins, and the other becomes a
/// no-op when it finds the continuation already consumed.
@MainActor
final class StreamCompletionWaiter {
private var continuation: CheckedContinuation<AppState.StreamCompletion?, Never>?

init(continuation: CheckedContinuation<AppState.StreamCompletion?, Never>) {
self.continuation = continuation
}

func resume(with value: AppState.StreamCompletion?) {
guard let cont = continuation else { return }
continuation = nil
cont.resume(returning: value)
}
}
6 changes: 1 addition & 5 deletions RxCode/App/AppState+Stream.swift
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,6 @@ extension AppState {
// New Claude turn after receiving tool result — start a new ChatMessage
state.messages[idx].isStreaming = false
state.messages[idx].finalizeToolCalls()
Self.stripNoOpText(at: idx, in: &state.messages)
state.needsNewMessage = false
state.messages.append(ChatMessage(role: .assistant, content: buffered, isStreaming: true))
} else {
Expand Down Expand Up @@ -305,7 +304,6 @@ extension AppState {
if let idx = state.messages.indices.reversed().first(where: { state.messages[$0].role == .assistant && state.messages[$0].isStreaming }) {
state.messages[idx].isStreaming = false
state.messages[idx].finalizeToolCalls()
Self.stripNoOpText(at: idx, in: &state.messages)
}
state.messages.append(ChatMessage(role: .assistant, isStreaming: true))
state.needsNewMessage = false
Expand Down Expand Up @@ -444,13 +442,11 @@ extension AppState {
// The user paused this turn — keep the partial assistant bubble
// visible. markStreamInterrupted() clears the message's streaming
// flag and retains in-progress tool calls (flagged as interrupted)
// instead of dropping them; the no-op strip below is told not to
// delete an emptied message.
// instead of dropping them.
state.messages[idx].markStreamInterrupted()
if let start = state.streamingStartDate {
state.messages[idx].duration = Date().timeIntervalSince(start)
}
Self.stripNoOpText(at: idx, in: &state.messages, removeIfEmpty: false)
}
state.streamingStartDate = nil
state.inFlightUserAttachments = []
Expand Down
12 changes: 12 additions & 0 deletions RxCode/App/AppState.swift
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,12 @@
/// id after the swap that happens mid-stream in `processStream`.
var sessionIdRedirect: [String: String] = [:]

/// Callers waiting for `pending-<streamId>` to be replaced by the CLI's
/// real `session_id` on the first `.system` event. The cross-project MCP
/// send (`ide__send_to_thread`) holds the JSON-RPC reply until the real
/// id is known so the sender's agent never sees a `pending-…` thread id.
var sessionIdRenameWaiters: [String: SessionRenameWaiter] = [:]

// MARK: - Stream Completion Tracking (cross-project MCP)

/// Result of a finished stream. Used by `ide__send_to_thread` to surface
Expand All @@ -186,6 +192,12 @@
/// have not yet picked up the result. Keyed by `streamId`.
var pendingStreamCompletions: [UUID: StreamCompletion] = [:]

/// Active callers waiting for a stream's completion. Resumed directly by
/// `recordStreamCompletion` so the cross-project MCP handler doesn't sit
/// in a polling loop on MainActor (which starves the target project's
/// `processStream` and freezes its thread until the sender is cancelled).
var streamCompletionWaiters: [UUID: StreamCompletionWaiter] = [:]

// MARK: - Session Summaries (shared — lightweight metadata for all projects)

var allSessionSummaries: [ChatSession.Summary] = []
Expand Down Expand Up @@ -999,4 +1011,4 @@
return message
}
}
}

Check warning on line 1014 in RxCode/App/AppState.swift

View workflow job for this annotation

GitHub Actions / swiftlint

File should contain 600 lines or less excluding comments and whitespaces: currently contains 664 (file_length)
Loading
Loading