Skip to content

Commit 2080c44

Browse files
OgeonX-AiAitomates
OgeonX-Ai
and
Aitomates
authored
fix(recovery): preserve failed MCP state for fault injection (#20)
Co-authored-by: Kim Harjamäki <kim.harjamaki@prosimo.fi>
1 parent a8c09ec commit 2080c44

3 files changed

Lines changed: 274 additions & 1 deletion

File tree

‎src/GsdOrchestrator.Tests/CheckpointCorruptionTests.cs‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
using System.Text.Json;
2+
using GsdOrchestrator.Checkpointing;
3+
using GsdOrchestrator.Workflows.Models;
4+
using Microsoft.Extensions.Logging.Abstractions;
5+
using Xunit;
6+
7+
namespace GsdOrchestrator.Tests;
8+
9+
public sealed class CheckpointCorruptionTests
10+
{
11+
[Fact]
12+
public async Task LoadAsync_CorruptedFinalCheckpoint_ThrowsJsonException()
13+
{
14+
using var fixture = new CheckpointFixture();
15+
var store = fixture.CreateStore();
16+
var context = ValidContext("corrupt-final");
17+
var path = fixture.StatePath(context.WorkflowId);
18+
19+
await store.SaveAsync(context);
20+
await File.WriteAllTextAsync(path, "{\"workflowId\":\"corrupt-final\",\"currentSt");
21+
22+
await Assert.ThrowsAsync<JsonException>(() => store.LoadAsync(context.WorkflowId));
23+
}
24+
25+
[Fact]
26+
public async Task LoadAsync_OrphanedTmpFile_IgnoresGarbageAndReturnsLastGoodCheckpoint()
27+
{
28+
using var fixture = new CheckpointFixture();
29+
var store = fixture.CreateStore();
30+
var context = ValidContext("orphaned-tmp");
31+
var path = fixture.StatePath(context.WorkflowId);
32+
var tmpPath = path + ".tmp";
33+
34+
await store.SaveAsync(context);
35+
await File.WriteAllTextAsync(tmpPath, "garbage bytes from interrupted second save");
36+
37+
var loaded = await store.LoadAsync(context.WorkflowId);
38+
39+
Assert.NotNull(loaded);
40+
Assert.Equal(context.WorkflowId, loaded!.WorkflowId);
41+
Assert.Equal(context.CurrentState, loaded.CurrentState);
42+
Assert.Equal("1.1", loaded.SchemaVersion);
43+
}
44+
45+
[Fact]
46+
public async Task LoadAsync_UnsupportedSchemaVersion_ThrowsInvalidDataException()
47+
{
48+
using var fixture = new CheckpointFixture();
49+
var store = fixture.CreateStore();
50+
var context = ValidContext("unsupported-version");
51+
52+
await store.SaveAsync(context);
53+
await RewriteSchemaVersionAsync(fixture.StatePath(context.WorkflowId), "99.0");
54+
55+
var error = await Assert.ThrowsAsync<InvalidDataException>(() => store.LoadAsync(context.WorkflowId));
56+
Assert.Contains("unsupported schema version", error.Message, StringComparison.OrdinalIgnoreCase);
57+
}
58+
59+
[Fact]
60+
public async Task LoadAsync_LegacySchemaVersion_UpgradesToCurrentVersion()
61+
{
62+
using var fixture = new CheckpointFixture();
63+
var store = fixture.CreateStore();
64+
var context = ValidContext("legacy-version");
65+
66+
await store.SaveAsync(context);
67+
await RewriteSchemaVersionAsync(fixture.StatePath(context.WorkflowId), "1.0");
68+
69+
var loaded = await store.LoadAsync(context.WorkflowId);
70+
71+
Assert.NotNull(loaded);
72+
Assert.Equal("1.1", loaded!.SchemaVersion);
73+
Assert.Equal(context.WorkflowId, loaded.WorkflowId);
74+
}
75+
76+
private static async Task RewriteSchemaVersionAsync(string path, string schemaVersion)
77+
{
78+
var checkpoint = await File.ReadAllTextAsync(path);
79+
var rewritten = checkpoint.Replace("\"schemaVersion\": \"1.1\"", $"\"schemaVersion\": \"{schemaVersion}\"", StringComparison.Ordinal);
80+
await File.WriteAllTextAsync(path, rewritten);
81+
}
82+
83+
private static GsdWorkflowContext ValidContext(string workflowId) =>
84+
new()
85+
{
86+
WorkflowId = workflowId,
87+
SchemaVersion = "1.1",
88+
CurrentState = WorkflowState.Analyzing,
89+
Issue = null
90+
};
91+
92+
private sealed class CheckpointFixture : IDisposable
93+
{
94+
private readonly string _root = Directory.CreateTempSubdirectory("gsd-checkpoint-tests-").FullName;
95+
96+
public FileCheckpointStore CreateStore() =>
97+
new(_root, NullLogger<FileCheckpointStore>.Instance);
98+
99+
public string StatePath(string workflowId) =>
100+
Path.Combine(_root, ".gsd", "state", $"{workflowId}.json");
101+
102+
public void Dispose()
103+
{
104+
if (Directory.Exists(_root))
105+
{
106+
Directory.Delete(_root, recursive: true);
107+
}
108+
}
109+
}
110+
}

‎src/GsdOrchestrator.Tests/FaultInjectionTests.cs‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
using GsdOrchestrator.Checkpointing;
2+
using GsdOrchestrator.Mcp;
3+
using GsdOrchestrator.Workflows;
4+
using GsdOrchestrator.Workflows.Models;
5+
using GsdOrchestrator.Workflows.States;
6+
using Microsoft.Extensions.Logging.Abstractions;
7+
using NSubstitute;
8+
using NSubstitute.ExceptionExtensions;
9+
using Polly.Registry;
10+
using Xunit;
11+
12+
namespace GsdOrchestrator.Tests;
13+
14+
public sealed class FaultInjectionTests
15+
{
16+
private static GsdStateMachine BuildSut(
17+
ICheckpointStore checkpoints,
18+
IWorkflowState[] states,
19+
IMcpClient? mcpClient = null)
20+
{
21+
var client = mcpClient ?? Substitute.For<IMcpClient>();
22+
var registry = new ResiliencePipelineRegistry<string>();
23+
registry.TryAddBuilder("mcp-tools", (b, _) => { });
24+
var dispatcher = new McpToolDispatcher(
25+
client,
26+
registry,
27+
NullLogger<McpToolDispatcher>.Instance);
28+
return new GsdStateMachine(
29+
checkpoints,
30+
dispatcher,
31+
states,
32+
NullLogger<GsdStateMachine>.Instance);
33+
}
34+
35+
private static IWorkflowState MakeTransitionState(WorkflowState from, WorkflowState to)
36+
{
37+
var state = Substitute.For<IWorkflowState>();
38+
state.State.Returns(from);
39+
state.ExecuteAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
40+
.Returns(ci => Task.FromResult(ci.Arg<GsdWorkflowContext>().Transition(to)));
41+
return state;
42+
}
43+
44+
[Fact]
45+
public async Task RunAsync_McpFailureMidGoal_ProducesTypedFailedStateAndCheckpoint()
46+
{
47+
var checkpoints = Substitute.For<ICheckpointStore>();
48+
checkpoints.SaveAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
49+
.Returns(Task.CompletedTask);
50+
51+
var idle = MakeTransitionState(WorkflowState.Idle, WorkflowState.Analyzing);
52+
var analyzing = MakeTransitionState(WorkflowState.Analyzing, WorkflowState.Editing);
53+
var editing = Substitute.For<IWorkflowState>();
54+
editing.State.Returns(WorkflowState.Editing);
55+
editing.ExecuteAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
56+
.ThrowsAsync(new McpException("simulated MCP failure", isTransient: true));
57+
58+
var sut = BuildSut(checkpoints, [idle, analyzing, editing]);
59+
60+
var ctx = await sut.RunAsync("owner", "repo", 42, triageModeOnly: false, CancellationToken.None);
61+
62+
Assert.Equal(WorkflowState.Failed, ctx.CurrentState);
63+
Assert.Equal(WorkflowState.Editing, ctx.FailedState);
64+
Assert.Contains("simulated MCP failure", ctx.FailureReason);
65+
await checkpoints.Received().SaveAsync(
66+
Arg.Is<GsdWorkflowContext>(saved =>
67+
saved.CurrentState == WorkflowState.Failed &&
68+
saved.FailedState == WorkflowState.Editing &&
69+
saved.FailureReason != null &&
70+
saved.FailureReason.Contains("simulated MCP failure", StringComparison.Ordinal)),
71+
Arg.Any<CancellationToken>());
72+
}
73+
74+
[Fact]
75+
public async Task ResumeAsync_RetriesFailedStateOnce_ThenHaltsDeterministically()
76+
{
77+
var checkpoints = Substitute.For<ICheckpointStore>();
78+
checkpoints.SaveAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
79+
.Returns(Task.CompletedTask);
80+
81+
var firstCheckpoint = FailedCheckpoint("retry-wf", retryCount: 0);
82+
var secondCheckpoint = FailedCheckpoint("retry-wf", retryCount: 1);
83+
checkpoints.LoadAsync("retry-wf", Arg.Any<CancellationToken>())
84+
.Returns(firstCheckpoint, secondCheckpoint);
85+
86+
var editing = Substitute.For<IWorkflowState>();
87+
editing.State.Returns(WorkflowState.Editing);
88+
editing.ExecuteAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
89+
.ThrowsAsync(new McpException("editing exploded again", isTransient: true));
90+
91+
var sut = BuildSut(checkpoints, [editing]);
92+
93+
var firstResult = await sut.ResumeAsync("retry-wf", CancellationToken.None);
94+
95+
Assert.Equal(WorkflowState.Failed, firstResult.CurrentState);
96+
Assert.Equal(1, firstResult.RetryCount);
97+
Assert.Equal(WorkflowState.Editing, firstResult.FailedState);
98+
Assert.Contains("editing exploded again", firstResult.FailureReason);
99+
await checkpoints.Received().SaveAsync(
100+
Arg.Is<GsdWorkflowContext>(saved =>
101+
saved.CurrentState == WorkflowState.Failed &&
102+
saved.RetryCount == 1 &&
103+
saved.FailedState == WorkflowState.Editing),
104+
Arg.Any<CancellationToken>());
105+
106+
var error = await Assert.ThrowsAsync<InvalidOperationException>(
107+
() => sut.ResumeAsync("retry-wf", CancellationToken.None));
108+
Assert.Contains("recovery retry limit", error.Message, StringComparison.OrdinalIgnoreCase);
109+
}
110+
111+
[Fact]
112+
public async Task RunAsync_TruncatesFailureReasonTo1024Characters()
113+
{
114+
var checkpoints = Substitute.For<ICheckpointStore>();
115+
checkpoints.SaveAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
116+
.Returns(Task.CompletedTask);
117+
118+
var oversizedMessage = new string('x', 1200);
119+
var idle = Substitute.For<IWorkflowState>();
120+
idle.State.Returns(WorkflowState.Idle);
121+
idle.ExecuteAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
122+
.ThrowsAsync(new InvalidOperationException(oversizedMessage));
123+
124+
var sut = BuildSut(checkpoints, [idle]);
125+
126+
var ctx = await sut.RunAsync("owner", "repo", 7, triageModeOnly: false, CancellationToken.None);
127+
128+
Assert.Equal(WorkflowState.Failed, ctx.CurrentState);
129+
Assert.NotNull(ctx.FailureReason);
130+
Assert.Equal(1024, ctx.FailureReason!.Length);
131+
}
132+
133+
[Fact]
134+
public async Task RunAsync_BudgetFailureMessage_SetsBudgetExhaustedStopReason()
135+
{
136+
var checkpoints = Substitute.For<ICheckpointStore>();
137+
checkpoints.SaveAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
138+
.Returns(Task.CompletedTask);
139+
140+
var idle = Substitute.For<IWorkflowState>();
141+
idle.State.Returns(WorkflowState.Idle);
142+
idle.ExecuteAsync(Arg.Any<GsdWorkflowContext>(), Arg.Any<CancellationToken>())
143+
.ThrowsAsync(new InvalidOperationException("budget exhausted during planning"));
144+
145+
var sut = BuildSut(checkpoints, [idle]);
146+
147+
var ctx = await sut.RunAsync("owner", "repo", 8, triageModeOnly: false, CancellationToken.None);
148+
149+
Assert.Equal(WorkflowState.Failed, ctx.CurrentState);
150+
Assert.Equal(TerminalStopReason.BudgetExhausted, ctx.StopReason);
151+
}
152+
153+
private static GsdWorkflowContext FailedCheckpoint(string workflowId, int retryCount) =>
154+
new()
155+
{
156+
WorkflowId = workflowId,
157+
CurrentState = WorkflowState.Failed,
158+
FailedState = WorkflowState.Editing,
159+
RetryCount = retryCount,
160+
FailureReason = "prior failure",
161+
Issue = new IssueContext(1, "Retry workflow", "", [], "owner", "repo", "main")
162+
};
163+
}

‎src/GsdOrchestrator/Workflows/GsdStateMachine.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,7 +235,7 @@ private async Task<GsdWorkflowContext> ExecuteLoopAsync(GsdWorkflowContext ctx,
235235
InvalidatedPhaseIds: [phase.RollbackTo ?? phase.Id],
236236
Reason: failureReason));
237237
}
238-
if (ctx.PendingRollback is not null)
238+
if (ctx.PendingRollback is not null && ex is not McpException)
239239
{
240240
_logger.LogWarning(
241241
"Workflow {WorkflowId} rollback requested from {FromPhase} to {ToPhase}: {Reason}",

0 commit comments

Comments
 (0)