Skip to content

Commit 64b49fb

Browse files
baronbaron
authored
Test MCP stdio and terminal commit barriers (#912)
1 parent 44b8205 commit 64b49fb

2 files changed

Lines changed: 207 additions & 0 deletions

File tree

‎Tests/RepoPromptTests/AgentMode/AgentRunLifecycleContractsTests.swift‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,4 +80,86 @@ final class AgentRunLifecycleContractsTests: XCTestCase {
8080
XCTAssertNil(contextBuilderSession.activeRunOwnership)
8181
XCTAssertTrue(contextBuilderSession.agentLog.isEmpty)
8282
}
83+
84+
@MainActor
85+
func testTerminalCommitRejectsStaleDrainAndDoesNotRepublishResolvedRevision() async throws {
86+
let tabID = UUID()
87+
let lifecycle = AgentRunAttemptLifecycle()
88+
let ownership = lifecycle.beginAttempt(
89+
context: .init(tabID: tabID, persistentSessionID: nil)
90+
)
91+
let providerDrainGeneration: UInt64 = 7
92+
var publicationCount = 0
93+
let hooks = AgentRunTerminalSessionBinding.Hooks(
94+
flushPendingAssistantDelta: {},
95+
finalizeStreamingItems: {},
96+
finalizePendingToolCalls: { _ in },
97+
finalizeNonCodexTurnUsage: {},
98+
cancelPendingInteractions: { _ in },
99+
finalizeAttachments: { _, _ in },
100+
setAgentRunInactive: {},
101+
prepareTerminalPublication: {},
102+
makeTerminalPublicationEnvelope: { _, _, _, _ in nil },
103+
updateBindings: {},
104+
notifyAgentTurnComplete: {},
105+
scheduleSave: {},
106+
publishTerminalCommit: { _, _ in
107+
publicationCount += 1
108+
return .accepted(successorEpoch: nil)
109+
},
110+
startFollowUpRun: { _ in }
111+
)
112+
let binding = AgentRunTerminalSessionBinding(
113+
tabID: tabID,
114+
lifecycle: lifecycle,
115+
hooks: hooks,
116+
validatesOwnership: { candidate, expectedRunID in
117+
lifecycle.isCurrentAttempt(candidate, expectedRunID: expectedRunID)
118+
},
119+
providerDrainGeneration: { providerDrainGeneration },
120+
terminalTurnID: { nil },
121+
queuedFollowUp: { nil },
122+
setFollowUpPending: { _ in },
123+
removeFirstQueuedFollowUp: { nil },
124+
appendError: { _ in },
125+
finishActiveState: { candidate, _, _ in
126+
_ = lifecycle.endAttempt(ifCurrent: candidate)
127+
},
128+
retainProcessRunIdentity: { _, _ in },
129+
sourceItemsRevision: { 3 },
130+
assistantDeltaFlushGeneration: { 5 },
131+
latestFailureText: { nil }
132+
)
133+
let barrier = AgentRunTerminalCommitBarrier()
134+
135+
func request(drainGeneration: UInt64) -> AgentRunTerminalCommitBarrier.Request {
136+
AgentRunTerminalCommitBarrier.Request(
137+
binding: binding,
138+
ownership: ownership,
139+
expectedRunID: nil,
140+
terminalState: .completed,
141+
source: "test.terminalCommit",
142+
attachmentDisposition: .deleteFiles,
143+
finalizeNonCodexUsage: false,
144+
supportsFollowUp: false,
145+
notifyTurnComplete: false,
146+
providerDrainGeneration: drainGeneration
147+
)
148+
}
149+
150+
let staleRevision = await barrier.commit(request(drainGeneration: providerDrainGeneration - 1))
151+
XCTAssertNil(staleRevision)
152+
XCTAssertEqual(publicationCount, 0)
153+
XCTAssertNil(lifecycle.lastTerminalCommitRevision)
154+
155+
let firstCandidate = await barrier.commit(request(drainGeneration: providerDrainGeneration))
156+
let firstRevision = try XCTUnwrap(firstCandidate)
157+
XCTAssertEqual(publicationCount, 1)
158+
XCTAssertEqual(lifecycle.lastTerminalPublicationResult, .accepted(successorEpoch: nil))
159+
160+
let repeatedCandidate = await barrier.commit(request(drainGeneration: providerDrainGeneration))
161+
let repeatedRevision = try XCTUnwrap(repeatedCandidate)
162+
XCTAssertEqual(repeatedRevision, firstRevision)
163+
XCTAssertEqual(publicationCount, 1)
164+
}
83165
}

‎Tests/RepoPromptTests/MCP/MCPStdioServerTransportTests.swift‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
import Darwin
2+
import Foundation
3+
@testable import RepoPromptMCP
4+
import XCTest
5+
6+
final class MCPStdioServerTransportTests: XCTestCase {
7+
func testCompleteFramesAreDeliveredExactlyOnceInOrderBeforeCleanEOF() async throws {
8+
var inputDescriptors = try makePipe()
9+
defer {
10+
closeDescriptor(&inputDescriptors[0])
11+
closeDescriptor(&inputDescriptors[1])
12+
}
13+
var outputDescriptors = try makePipe()
14+
defer {
15+
closeDescriptor(&outputDescriptors[0])
16+
closeDescriptor(&outputDescriptors[1])
17+
}
18+
19+
let transport = MCPStdioServerTransport(
20+
stdinFD: inputDescriptors[0],
21+
stdoutFD: outputDescriptors[1],
22+
parentPIDProvider: { 42 }
23+
)
24+
try await transport.connect()
25+
let stream = await transport.receive()
26+
async let terminal = transport.waitUntilTerminal()
27+
28+
let expectedFrames = [
29+
Data(#"{"jsonrpc":"2.0","id":1}"#.utf8),
30+
Data(#"{"jsonrpc":"2.0","id":2}"#.utf8)
31+
]
32+
var wireBytes = Data()
33+
for frame in expectedFrames {
34+
wireBytes.append(frame)
35+
wireBytes.append(0x0A)
36+
}
37+
writeExactly(wireBytes, to: inputDescriptors[1])
38+
closeDescriptor(&inputDescriptors[1])
39+
40+
var receivedFrames: [Data] = []
41+
for try await frame in stream {
42+
receivedFrames.append(frame)
43+
}
44+
45+
let observedTerminal = await terminal
46+
XCTAssertEqual(receivedFrames, expectedFrames)
47+
XCTAssertEqual(observedTerminal, .stdinEOF)
48+
await transport.disconnect()
49+
}
50+
51+
func testEOFWithIncompleteFrameReportsExactTruncatedByteCount() async throws {
52+
var inputDescriptors = try makePipe()
53+
defer {
54+
closeDescriptor(&inputDescriptors[0])
55+
closeDescriptor(&inputDescriptors[1])
56+
}
57+
var outputDescriptors = try makePipe()
58+
defer {
59+
closeDescriptor(&outputDescriptors[0])
60+
closeDescriptor(&outputDescriptors[1])
61+
}
62+
63+
let transport = MCPStdioServerTransport(
64+
stdinFD: inputDescriptors[0],
65+
stdoutFD: outputDescriptors[1],
66+
parentPIDProvider: { 42 }
67+
)
68+
try await transport.connect()
69+
let stream = await transport.receive()
70+
async let terminal = transport.waitUntilTerminal()
71+
72+
let incompleteFrame = Data(#"{"jsonrpc":"2.0""#.utf8)
73+
writeExactly(incompleteFrame, to: inputDescriptors[1])
74+
closeDescriptor(&inputDescriptors[1])
75+
76+
var receivedFrames: [Data] = []
77+
var streamError: Error?
78+
do {
79+
for try await frame in stream {
80+
receivedFrames.append(frame)
81+
}
82+
} catch {
83+
streamError = error
84+
}
85+
86+
let expectedTerminal = MCPStdioServerTransport.TerminalError.stdinTruncatedFrame(
87+
bytes: incompleteFrame.count
88+
)
89+
let observedTerminal = await terminal
90+
XCTAssertTrue(receivedFrames.isEmpty)
91+
XCTAssertEqual(streamError as? MCPStdioServerTransport.TerminalError, expectedTerminal)
92+
XCTAssertEqual(observedTerminal, expectedTerminal)
93+
await transport.disconnect()
94+
}
95+
96+
private func makePipe() throws -> [Int32] {
97+
var descriptors = [Int32](repeating: -1, count: 2)
98+
guard Darwin.pipe(&descriptors) == 0 else {
99+
throw NSError(domain: NSPOSIXErrorDomain, code: Int(errno))
100+
}
101+
return descriptors
102+
}
103+
104+
private func writeExactly(
105+
_ data: Data,
106+
to descriptor: Int32,
107+
file: StaticString = #filePath,
108+
line: UInt = #line
109+
) {
110+
let written = data.withUnsafeBytes { bytes in
111+
Darwin.write(descriptor, bytes.baseAddress, bytes.count)
112+
}
113+
XCTAssertEqual(written, data.count, file: file, line: line)
114+
}
115+
116+
private func closeDescriptor(
117+
_ descriptor: inout Int32,
118+
file: StaticString = #filePath,
119+
line: UInt = #line
120+
) {
121+
guard descriptor >= 0 else { return }
122+
XCTAssertEqual(Darwin.close(descriptor), 0, file: file, line: line)
123+
descriptor = -1
124+
}
125+
}

0 commit comments

Comments
 (0)