@@ -7,6 +7,7 @@ import { validateSnapshot } from './schemaValidation.js'
7
7
import { pingIntervalId } from './utils.js' ;
8
8
import { startPolling } from './utils.js' ;
9
9
10
+ const uploadDomToS3ViaEnv = process . env . USE_LAMBDA_INTERNAL || false ;
10
11
export default async ( ctx : Context ) : Promise < FastifyInstance < Server , IncomingMessage , ServerResponse > > => {
11
12
12
13
const server : FastifyInstance < Server , IncomingMessage , ServerResponse > = fastify ( {
@@ -41,6 +42,7 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
41
42
// Fetch sessionId from snapshot options if present
42
43
const sessionId = snapshot ?. options ?. sessionId ;
43
44
let capsBuildId = ''
45
+ const contextId = snapshot ?. options ?. contextId ;
44
46
45
47
if ( sessionId ) {
46
48
// Check if sessionId exists in the map
@@ -71,7 +73,23 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
71
73
}
72
74
73
75
ctx . testType = testType ;
74
- ctx . snapshotQueue ?. enqueue ( snapshot ) ;
76
+
77
+ if ( contextId && ! ctx . contextToSnapshotMap ) {
78
+ ctx . contextToSnapshotMap = new Map ( ) ;
79
+ ctx . log . debug ( `Initialized empty context mapping map for contextId: ${ contextId } ` ) ;
80
+ }
81
+
82
+ if ( contextId && ctx . contextToSnapshotMap ) {
83
+ ctx . contextToSnapshotMap . set ( contextId , 0 ) ;
84
+ ctx . log . debug ( `Marking contextId as captured and added to queue: ${ contextId } ` ) ;
85
+ }
86
+
87
+ if ( contextId ) {
88
+ ctx . snapshotQueue ?. enqueueFront ( snapshot ) ;
89
+ } else {
90
+ ctx . snapshotQueue ?. enqueue ( snapshot ) ;
91
+ }
92
+
75
93
ctx . isSnapshotCaptured = true ;
76
94
replyCode = 200 ;
77
95
replyBody = { data : { message : "success" , warnings : [ ] } } ;
@@ -105,8 +123,16 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
105
123
if ( ctx . server ) {
106
124
ctx . server . close ( ) ;
107
125
}
108
- let resp = await ctx . client . getS3PreSignedURL ( ctx ) ;
109
- await ctx . client . uploadLogs ( ctx , resp . data . url ) ;
126
+
127
+ let uploadCLILogsToS3 = ctx ?. config ?. useLambdaInternal || uploadDomToS3ViaEnv ;
128
+ if ( ! uploadCLILogsToS3 ) {
129
+ ctx . log . debug ( `Log file to be uploaded` )
130
+ let resp = await ctx . client . getS3PreSignedURL ( ctx ) ;
131
+ await ctx . client . uploadLogs ( ctx , resp . data . url ) ;
132
+ } else {
133
+ ctx . log . debug ( `Log file to be uploaded via LSRS` )
134
+ let resp = ctx . client . sendCliLogsToLSRS ( ctx ) ;
135
+ }
110
136
111
137
if ( pingIntervalId !== null ) {
112
138
clearInterval ( pingIntervalId ) ;
@@ -130,7 +156,113 @@ export default async (ctx: Context): Promise<FastifyInstance<Server, IncomingMes
130
156
reply . code ( 200 ) . send ( { status : 'Server is running' , version : ctx . cliVersion } ) ;
131
157
} ) ;
132
158
133
-
159
+ // Get snapshot status
160
+ server . get ( '/snapshot/status' , opts , async ( request , reply ) => {
161
+ let replyCode : number ;
162
+ let replyBody : Record < string , any > ;
163
+
164
+
165
+ try {
166
+ ctx . log . debug ( `request.query : ${ JSON . stringify ( request . query ) } ` ) ;
167
+ const { contextId, pollTimeout, snapshotName } = request . query as { contextId : string , pollTimeout : number , snapshotName : string } ;
168
+ if ( ! contextId || ! snapshotName ) {
169
+ throw new Error ( 'contextId and snapshotName are required parameters' ) ;
170
+ }
171
+
172
+ const timeoutDuration = pollTimeout * 1000 || 30000 ;
173
+
174
+ // Check if we have stored snapshot status for this contextId
175
+ if ( ctx . contextToSnapshotMap ?. has ( contextId ) ) {
176
+ let contextStatus = ctx . contextToSnapshotMap . get ( contextId ) ;
177
+
178
+ while ( contextStatus == 0 ) {
179
+ // Wait 5 seconds before next check
180
+ await new Promise ( resolve => setTimeout ( resolve , 5000 ) ) ;
181
+
182
+ contextStatus = ctx . contextToSnapshotMap . get ( contextId ) ;
183
+ }
184
+
185
+ if ( contextStatus == 2 ) {
186
+ throw new Error ( "Snapshot Failed" ) ;
187
+ }
188
+
189
+ ctx . log . debug ( "Snapshot uploaded successfully" ) ;
190
+
191
+ // Poll external API until it returns 200 or timeout is reached
192
+ let lastExternalResponse : any = null ;
193
+ const startTime = Date . now ( ) ;
194
+
195
+ while ( true ) {
196
+ try {
197
+ const externalResponse = await ctx . client . getSnapshotStatus (
198
+ snapshotName ,
199
+ contextId ,
200
+ ctx
201
+ ) ;
202
+
203
+ lastExternalResponse = externalResponse ;
204
+
205
+ if ( externalResponse . statusCode === 200 ) {
206
+ replyCode = 200 ;
207
+ replyBody = externalResponse . data ;
208
+ return reply . code ( replyCode ) . send ( replyBody ) ;
209
+ } else if ( externalResponse . statusCode === 202 ) {
210
+ replyBody = externalResponse . data ;
211
+ ctx . log . debug ( `External API attempt: Still processing, Pending Screenshots ${ externalResponse . snapshotCount } ` ) ;
212
+ await new Promise ( resolve => setTimeout ( resolve , 5000 ) ) ;
213
+ } else if ( externalResponse . statusCode === 404 ) {
214
+ ctx . log . debug ( `Snapshot still processing, not uploaded` ) ;
215
+ await new Promise ( resolve => setTimeout ( resolve , 5000 ) ) ;
216
+ } else {
217
+ ctx . log . debug ( `Unexpected response from external API: ${ JSON . stringify ( externalResponse ) } ` ) ;
218
+ replyCode = 500 ;
219
+ replyBody = {
220
+ error : {
221
+ message : `Unexpected response from external API: ${ externalResponse . statusCode } ` ,
222
+ externalApiStatus : externalResponse . statusCode
223
+ }
224
+ } ;
225
+ return reply . code ( replyCode ) . send ( replyBody ) ;
226
+ }
227
+
228
+ ctx . log . debug ( `timeoutDuration: ${ timeoutDuration } ` ) ;
229
+ ctx . log . debug ( `Time passed: ${ Date . now ( ) - startTime } ` ) ;
230
+
231
+ if ( Date . now ( ) - startTime > timeoutDuration ) {
232
+ replyCode = 202 ;
233
+ replyBody = {
234
+ data : {
235
+ message : 'Request timed out-> Snapshot still processing'
236
+ }
237
+ } ;
238
+ return reply . code ( replyCode ) . send ( replyBody ) ;
239
+ }
240
+
241
+ } catch ( externalApiError : any ) {
242
+ ctx . log . debug ( `External API call failed: ${ externalApiError . message } ` ) ;
243
+ replyCode = 500 ;
244
+ replyBody = {
245
+ error : {
246
+ message : `External API call failed: ${ externalApiError . message } `
247
+ }
248
+ } ;
249
+ return reply . code ( replyCode ) . send ( replyBody ) ;
250
+ }
251
+ }
252
+ } else {
253
+ // No snapshot found for this contextId
254
+ replyCode = 404 ;
255
+ replyBody = { error : { message : `No snapshot found for contextId: ${ contextId } ` } } ;
256
+ return reply . code ( replyCode ) . send ( replyBody ) ;
257
+ }
258
+ } catch ( error : any ) {
259
+ ctx . log . debug ( `snapshot status failed; ${ error } ` ) ;
260
+ replyCode = 500 ;
261
+ replyBody = { error : { message : error . message } } ;
262
+ return reply . code ( replyCode ) . send ( replyBody ) ;
263
+ }
264
+ } ) ;
265
+
134
266
135
267
await server . listen ( { port : ctx . options . port } ) ;
136
268
// store server's address for SDK
0 commit comments