copilot-sdk/dotnet/test/SystemMessageTransformTests.cs at main · github/copilot-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *--------------------------------------------------------------------------------------------*/

using GitHub.Copilot.SDK.Test.Harness;
using Xunit;
using Xunit.Abstractions;

namespace GitHub.Copilot.SDK.Test;

public class SystemMessageTransformTests(E2ETestFixture fixture, ITestOutputHelper output) : E2ETestBase(fixture, "system_message_transform", output)
{
    [Fact]
    public async Task Should_Invoke_Transform_Callbacks_With_Section_Content()
    {
        var identityCallbackInvoked = false;
        var toneCallbackInvoked = false;

        var session = await CreateSessionAsync(new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
            SystemMessage = new SystemMessageConfig
            {
                Mode = SystemMessageMode.Customize,
                Sections = new Dictionary<string, SectionOverride>
                {
                    ["identity"] = new SectionOverride
                    {
                        Transform = async (content) =>
                        {
                            Assert.False(string.IsNullOrEmpty(content));
                            identityCallbackInvoked = true;
                            return content;
                        }
                    },
                    ["tone"] = new SectionOverride
                    {
                        Transform = async (content) =>
                        {
                            Assert.False(string.IsNullOrEmpty(content));
                            toneCallbackInvoked = true;
                            return content;
                        }
                    }
                }
            }
        });

        await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "test.txt"), "Hello transform!");

        await session.SendAsync(new MessageOptions
        {
            Prompt = "Read the contents of test.txt and tell me what it says"
        });

        await TestHelper.GetFinalAssistantMessageAsync(session);

        Assert.True(identityCallbackInvoked, "Expected identity transform callback to be invoked");
        Assert.True(toneCallbackInvoked, "Expected tone transform callback to be invoked");
    }

    [Fact]
    public async Task Should_Apply_Transform_Modifications_To_Section_Content()
    {
        var session = await CreateSessionAsync(new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
            SystemMessage = new SystemMessageConfig
            {
                Mode = SystemMessageMode.Customize,
                Sections = new Dictionary<string, SectionOverride>
                {
                    ["identity"] = new SectionOverride
                    {
                        Transform = async (content) =>
                        {
                            return content + "\nAlways end your reply with TRANSFORM_MARKER";
                        }
                    }
                }
            }
        });

        await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "hello.txt"), "Hello!");

        await session.SendAsync(new MessageOptions
        {
            Prompt = "Read the contents of hello.txt"
        });

        await TestHelper.GetFinalAssistantMessageAsync(session);

        // Verify the transform result was actually applied to the system message
        var traffic = await Ctx.GetExchangesAsync();
        Assert.NotEmpty(traffic);
        var systemMessage = GetSystemMessage(traffic[0]);
        Assert.Contains("TRANSFORM_MARKER", systemMessage);
    }

    [Fact]
    public async Task Should_Work_With_Static_Overrides_And_Transforms_Together()
    {
        var transformCallbackInvoked = false;

        var session = await CreateSessionAsync(new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
            SystemMessage = new SystemMessageConfig
            {
                Mode = SystemMessageMode.Customize,
                Sections = new Dictionary<string, SectionOverride>
                {
                    ["safety"] = new SectionOverride
                    {
                        Action = SectionOverrideAction.Remove
                    },
                    ["identity"] = new SectionOverride
                    {
                        Transform = async (content) =>
                        {
                            transformCallbackInvoked = true;
                            return content;
                        }
                    }
                }
            }
        });

        await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "combo.txt"), "Combo test!");

        await session.SendAsync(new MessageOptions
        {
            Prompt = "Read the contents of combo.txt and tell me what it says"
        });

        await TestHelper.GetFinalAssistantMessageAsync(session);

        Assert.True(transformCallbackInvoked, "Expected identity transform callback to be invoked");
    }
}