1
// Copyright (c) 2024 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
import * as assert from 'node:assert';
6
 
7
export const tests = {
8
  async test(_, env) {
9
    {
10
      // Test gateway get url
11
      const resp = await env.ai.gateway('my-gateway').getUrl('openai');
12
      assert.deepEqual(
13
        resp,
14
        'https://gateway.ai.cloudflare.com/v1/account-tag-abc/my-gateway/openai'
15
      );
16
    }
17
 
18
    {
19
      // Test gateway get log
20
      const resp = await env.ai.gateway('my-gateway').getLog('my-log-123');
21
      assert.deepEqual(resp, {
22
        cached: false,
23
        cost: 0,
24
        created_at: new Date('2019-08-24T14:15:22Z'),
25
        custom_cost: true,
26
        duration: 0,
27
        id: 'string',
28
        metadata: 'string',
29
        model: 'string',
30
        model_type: 'string',
31
        path: 'string',
32
        provider: 'string',
33
        request_content_type: 'string',
34
        request_head: 'string',
35
        request_head_complete: true,
36
        request_size: 0,
37
        request_type: 'string',
38
        response_content_type: 'string',
39
        response_head: 'string',
40
        response_head_complete: true,
41
        response_size: 0,
42
        status_code: 0,
43
        step: 0,
44
        success: true,
45
        tokens_in: 0,
46
        tokens_out: 0,
47
      });
48
    }
49
 
50
    {
51
      // Test get log error responses
52
      try {
53
        await env.ai.gateway('my-gateway').getLog('404');
54
      } catch (e) {
55
        assert.deepEqual(
56
          {
57
            name: e.name,
58
            message: e.message,
59
          },
60
          {
61
            name: 'AiGatewayLogNotFound',
62
            message: 'Not Found',
63
          }
64
        );
65
      }
66
    }
67
 
68
    {
69
      try {
70
        await env.ai.gateway('my-gateway').getLog('500');
71
      } catch (e) {
72
        assert.deepEqual(
73
          {
74
            name: e.name,
75
            message: e.message,
76
          },
77
          {
78
            name: 'AiGatewayInternalError',
79
            message: 'Internal Error',
80
          }
81
        );
82
      }
83
    }
84
 
85
    {
86
      // Test patch log error responses
87
      try {
88
        await env.ai.gateway('my-gateway').patchLog('404', { feedback: -1 });
89
      } catch (e) {
90
        assert.deepEqual(
91
          {
92
            name: e.name,
93
            message: e.message,
94
          },
95
          {
96
            name: 'AiGatewayLogNotFound',
97
            message: 'Not Found',
98
          }
99
        );
100
      }
101
    }
102
 
103
    {
104
      try {
105
        await env.ai.gateway('my-gateway').patchLog('500', { feedback: -1 });
106
      } catch (e) {
107
        assert.deepEqual(
108
          {
109
            name: e.name,
110
            message: e.message,
111
          },
112
          {
113
            name: 'AiGatewayInternalError',
114
            message: 'Internal Error',
115
          }
116
        );
117
      }
118
    }
119
 
120
    // Universal Run
121
    {
122
      const resp = await env.ai.gateway('my-gateway').run({
123
        provider: 'workers-ai',
124
        endpoint: '@cf/meta/llama-3.1-8b-instruct',
125
        headers: {
126
          Authorization: 'Bearer abcde',
127
          'Content-Type': 'application/json',
128
          'cf-aig-metadata': { user: 123 },
129
          'cf-aig-custom-cost': { total_cost: 1.22 },
130
          'cf-aig-skip-cache': false,
131
          'cf-aig-cache-ttl': 123,
132
        },
133
        query: {
134
          prompt: 'What is Cloudflare?',
135
        },
136
      });
137
 
138
      const body = await resp.json();
139
 
140
      assert.deepEqual(body, {
141
        headers: {
142
          'content-length': '322',
143
          'content-type': 'application/json',
144
        },
145
        result: [
146
          {
147
            endpoint: '@cf/meta/llama-3.1-8b-instruct',
148
            headers: {
149
              'Content-Type': 'application/json',
150
              'cf-aig-cache-ttl': '123',
151
              'cf-aig-custom-cost': '{"total_cost":1.22}',
152
              'cf-aig-metadata': '{"user":123}',
153
              'cf-aig-skip-cache': 'false',
154
              Authorization: 'Bearer abcde',
155
            },
156
            provider: 'workers-ai',
157
            query: { prompt: 'What is Cloudflare?' },
158
          },
159
        ],
160
        success: true,
161
      });
162
    }
163
 
164
    // Universal Run with global options
165
    {
166
      const resp = await env.ai.gateway('my-gateway').run(
167
        {
168
          provider: 'workers-ai',
169
          endpoint: '@cf/meta/llama-3.1-8b-instruct',
170
          headers: {
171
            Authorization: 'Bearer abcde',
172
            'Content-Type': 'application/json',
173
          },
174
          query: {
175
            prompt: 'What is Cloudflare?',
176
          },
177
        },
178
        {
179
          gateway: {
180
            skipCache: false,
181
            cacheTtl: 123,
182
          },
183
          extraHeaders: {
184
            'x-custom': 'stuff',
185
          },
186
        }
187
      );
188
 
189
      const body = await resp.json();
190
 
191
      assert.deepEqual(body, {
192
        headers: {
193
          'content-length': '189',
194
          'content-type': 'application/json',
195
          'cf-aig-skip-cache': 'false',
196
          'cf-aig-cache-ttl': '123',
197
          'x-custom': 'stuff',
198
        },
199
        result: [
200
          {
201
            endpoint: '@cf/meta/llama-3.1-8b-instruct',
202
            headers: {
203
              'Content-Type': 'application/json',
204
              Authorization: 'Bearer abcde',
205
            },
206
            provider: 'workers-ai',
207
            query: { prompt: 'What is Cloudflare?' },
208
          },
209
        ],
210
        success: true,
211
      });
212
    }
213
 
214
    {
215
      // Test already-aborted signal throws AbortError
216
      await assert.rejects(
217
        async () => {
218
          await env.ai.gateway('my-gateway').run(
219
            {
220
              provider: 'workers-ai',
221
              endpoint: '@cf/meta/llama-3.1-8b-instruct',
222
              headers: {
223
                Authorization: 'Bearer abcde',
224
                'Content-Type': 'application/json',
225
              },
226
              query: {
227
                prompt: 'What is Cloudflare?',
228
              },
229
            },
230
            { signal: AbortSignal.abort() }
231
          );
232
        },
233
        { name: 'AbortError' }
234
      );
235
    }
236
 
237
    {
238
      // Test aborting an in-flight request
239
      const controller = new AbortController();
240
      let resolved = false;
241
      const promise = env.ai
242
        .gateway('my-gateway')
243
        .run(
244
          {
245
            provider: 'workers-ai',
246
            endpoint: 'hang',
247
            headers: {
248
              Authorization: 'Bearer abcde',
249
              'Content-Type': 'application/json',
250
            },
251
            query: {
252
              prompt: 'What is Cloudflare?',
253
            },
254
          },
255
          { signal: controller.signal }
256
        )
257
        .finally(() => {
258
          resolved = true;
259
        });
260
      await scheduler.wait(10);
261
      assert.deepStrictEqual(resolved, false);
262
      controller.abort();
263
      await assert.rejects(promise, { name: 'AbortError' });
264
    }
265
  },
266
};