copilot-sdk/nodejs/test/e2e/system_message_transform.e2e.test.ts 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
/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *--------------------------------------------------------------------------------------------*/

import { writeFile } from "fs/promises";
import { join } from "path";
import { describe, expect, it } from "vitest";
import { ParsedHttpExchange } from "../../../test/harness/replayingCapiProxy.js";
import { approveAll } from "../../src/index.js";
import { createSdkTestContext } from "./harness/sdkTestContext.js";

describe("System message transform", async () => {
    const { copilotClient: client, openAiEndpoint, workDir } = await createSdkTestContext();

    it("should invoke transform callbacks with section content", async () => {
        const transformedSections: Record<string, string> = {};

        const session = await client.createSession({
            onPermissionRequest: approveAll,
            systemMessage: {
                mode: "customize",
                sections: {
                    identity: {
                        action: (content: string) => {
                            transformedSections["identity"] = content;
                            // Pass through unchanged
                            return content;
                        },
                    },
                    tone: {
                        action: (content: string) => {
                            transformedSections["tone"] = content;
                            return content;
                        },
                    },
                },
            },
        });

        await writeFile(join(workDir, "test.txt"), "Hello transform!");

        await session.sendAndWait({
            prompt: "Read the contents of test.txt and tell me what it says",
        });

        // Transform callbacks should have been invoked with real section content
        expect(Object.keys(transformedSections).length).toBe(2);
        expect(transformedSections["identity"]).toBeDefined();
        expect(transformedSections["identity"]!.length).toBeGreaterThan(0);
        expect(transformedSections["tone"]).toBeDefined();
        expect(transformedSections["tone"]!.length).toBeGreaterThan(0);

        await session.disconnect();
    });

    it("should apply transform modifications to section content", async () => {
        const session = await client.createSession({
            onPermissionRequest: approveAll,
            systemMessage: {
                mode: "customize",
                sections: {
                    identity: {
                        action: (content: string) => {
                            return content + "\nTRANSFORM_MARKER";
                        },
                    },
                },
            },
        });

        await writeFile(join(workDir, "hello.txt"), "Hello!");

        await session.sendAndWait({
            prompt: "Read the contents of hello.txt",
        });

        // Verify the transform result was actually applied to the system message
        const traffic = await openAiEndpoint.getExchanges();
        const systemMessage = getSystemMessage(traffic[0]);
        expect(systemMessage).toContain("TRANSFORM_MARKER");

        await session.disconnect();
    });

    it("should work with static overrides and transforms together", async () => {
        const transformedSections: Record<string, string> = {};

        const session = await client.createSession({
            onPermissionRequest: approveAll,
            systemMessage: {
                mode: "customize",
                sections: {
                    // Static override
                    safety: { action: "remove" },
                    // Transform
                    identity: {
                        action: (content: string) => {
                            transformedSections["identity"] = content;
                            return content;
                        },
                    },
                },
            },
        });

        await writeFile(join(workDir, "combo.txt"), "Combo test!");

        await session.sendAndWait({
            prompt: "Read the contents of combo.txt and tell me what it says",
        });

        // Transform should have been invoked
        expect(transformedSections["identity"]).toBeDefined();
        expect(transformedSections["identity"]!.length).toBeGreaterThan(0);

        await session.disconnect();
    });
});

function getSystemMessage(exchange: ParsedHttpExchange): string | undefined {
    const systemMessage = exchange.request.messages.find((m) => m.role === "system") as
        | { role: "system"; content: string }
        | undefined;
    return systemMessage?.content;
}