Skip to content

Commit 6ca8cb7

Browse files
ericallamericallam
committed
WIP
1 parent be02439 commit 6ca8cb7

File tree

6 files changed

+145
-68
lines changed

6 files changed

+145
-68
lines changed

‎apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.complete.ts

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,20 @@ import { logger } from "~/services/logger.server";
1313
import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
1414
import { engine } from "~/v3/runEngine.server";
1515

16-
const { action } = createActionApiRoute(
16+
const { action, loader } = createActionApiRoute(
1717
{
1818
params: z.object({
1919
waitpointFriendlyId: z.string(),
2020
}),
2121
body: CompleteWaitpointTokenRequestBody,
2222
maxContentLength: env.TASK_PAYLOAD_MAXIMUM_SIZE,
23-
method: "POST",
23+
allowJWT: true,
24+
authorization: {
25+
action: "write",
26+
resource: (params) => ({ waitpoints: params.waitpointFriendlyId }),
27+
superScopes: ["write:waitpoints", "admin"],
28+
},
29+
corsStrategy: "all",
2430
},
2531
async ({ authentication, body, params }) => {
2632
// Resume tokens are actually just waitpoints
@@ -65,4 +71,4 @@ const { action } = createActionApiRoute(
6571
}
6672
);
6773

68-
export { action };
74+
export { action, loader };

‎apps/webapp/app/routes/api.v1.waitpoints.tokens.ts

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
createActionApiRoute,
1414
createLoaderApiRoute,
1515
} from "~/services/routeBuilders/apiBuilder.server";
16+
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
1617
import { parseDelay } from "~/utils/delays";
1718
import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
1819
import { engine } from "~/v3/runEngine.server";
@@ -77,12 +78,14 @@ const { action } = createActionApiRoute(
7778
tags: bodyTags,
7879
});
7980

81+
const $responseHeaders = await responseHeaders(authentication.environment);
82+
8083
return json<CreateWaitpointTokenResponseBody>(
8184
{
8285
id: WaitpointId.toFriendlyId(result.waitpoint.id),
8386
isCached: result.isCached,
8487
},
85-
{ status: 200 }
88+
{ status: 200, headers: $responseHeaders }
8689
);
8790
} catch (error) {
8891
if (error instanceof ServiceValidationError) {
@@ -96,4 +99,17 @@ const { action } = createActionApiRoute(
9699
}
97100
);
98101

102+
async function responseHeaders(
103+
environment: AuthenticatedEnvironment
104+
): Promise<Record<string, string>> {
105+
const claimsHeader = JSON.stringify({
106+
sub: environment.id,
107+
pub: true,
108+
});
109+
110+
return {
111+
"x-trigger-jwt-claims": claimsHeader,
112+
};
113+
}
114+
99115
export { action };

‎apps/webapp/app/services/authorization.server.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
export type AuthorizationAction = "read" | "write" | string; // Add more actions as needed
22

3-
const ResourceTypes = ["tasks", "tags", "runs", "batch"] as const;
3+
const ResourceTypes = ["tasks", "tags", "runs", "batch", "waitpoints"] as const;
44

55
export type AuthorizationResources = {
66
[key in (typeof ResourceTypes)[number]]?: string | string[];

‎packages/core/src/v3/apiClient/core.ts

Lines changed: 34 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { calculateNextRetryDelay } from "../utils/retries.js";
55
import { ApiConnectionError, ApiError, ApiSchemaValidationError } from "./errors.js";
66

77
import { Attributes, context, propagation, Span } from "@opentelemetry/api";
8-
import {suppressTracing} from "@opentelemetry/core"
8+
import { suppressTracing } from "@opentelemetry/core";
99
import { SemanticInternalAttributes } from "../semanticInternalAttributes.js";
1010
import type { TriggerTracer } from "../tracer.js";
1111
import { accessoryAttributes } from "../utils/styleAttributes.js";
@@ -27,14 +27,14 @@ export const defaultRetryOptions = {
2727
randomize: false,
2828
} satisfies RetryOptions;
2929

30-
export type ZodFetchOptions<T = unknown> = {
30+
export type ZodFetchOptions<TInput = unknown, TOutput = TInput> = {
3131
retry?: RetryOptions;
3232
tracer?: TriggerTracer;
3333
name?: string;
3434
attributes?: Attributes;
3535
icon?: string;
36-
onResponseBody?: (body: T, span: Span) => void;
37-
prepareData?: (data: T) => Promise<T> | T;
36+
onResponseBody?: (body: TInput, span: Span) => void;
37+
prepareData?: (data: TInput, response: Response) => Promise<TOutput> | TOutput;
3838
};
3939

4040
export type AnyZodFetchOptions = ZodFetchOptions<any>;
@@ -67,12 +67,15 @@ interface FetchOffsetLimitPageParams extends OffsetLimitPageParams {
6767
query?: URLSearchParams;
6868
}
6969

70-
export function zodfetch<TResponseBodySchema extends z.ZodTypeAny>(
70+
export function zodfetch<
71+
TResponseBodySchema extends z.ZodTypeAny,
72+
TOutput = z.output<TResponseBodySchema>,
73+
>(
7174
schema: TResponseBodySchema,
7275
url: string,
7376
requestInit?: RequestInit,
74-
options?: ZodFetchOptions<z.output<TResponseBodySchema>>
75-
): ApiPromise<z.output<TResponseBodySchema>> {
77+
options?: ZodFetchOptions<z.output<TResponseBodySchema>, TOutput>
78+
): ApiPromise<TOutput> {
7679
return new ApiPromise(_doZodFetch(schema, url, requestInit, options));
7780
}
7881

@@ -110,7 +113,14 @@ export function zodfetchCursorPage<TItemSchema extends z.ZodTypeAny>(
110113

111114
const fetchResult = _doZodFetch(cursorPageSchema, $url.href, requestInit, options);
112115

113-
return new CursorPagePromise(fetchResult, schema, url, params, requestInit, options);
116+
return new CursorPagePromise(
117+
fetchResult as Promise<ZodFetchResult<CursorPageResponse<z.output<TItemSchema>>>>,
118+
schema,
119+
url,
120+
params,
121+
requestInit,
122+
options
123+
);
114124
}
115125

116126
export function zodfetchOffsetLimitPage<TItemSchema extends z.ZodTypeAny>(
@@ -144,7 +154,14 @@ export function zodfetchOffsetLimitPage<TItemSchema extends z.ZodTypeAny>(
144154

145155
const fetchResult = _doZodFetch(offsetLimitPageSchema, $url.href, requestInit, options);
146156

147-
return new OffsetLimitPagePromise(fetchResult, schema, url, params, requestInit, options);
157+
return new OffsetLimitPagePromise(
158+
fetchResult as Promise<ZodFetchResult<OffsetLimitPageResponse<z.output<TItemSchema>>>>,
159+
schema,
160+
url,
161+
params,
162+
requestInit,
163+
options
164+
);
148165
}
149166

150167
type ZodFetchResult<T> = {
@@ -184,12 +201,15 @@ async function traceZodFetch<T>(
184201
);
185202
}
186203

187-
async function _doZodFetch<TResponseBodySchema extends z.ZodTypeAny>(
204+
async function _doZodFetch<
205+
TResponseBodySchema extends z.ZodTypeAny,
206+
TOutput = z.output<TResponseBodySchema>,
207+
>(
188208
schema: TResponseBodySchema,
189209
url: string,
190210
requestInit?: PromiseOrValue<RequestInit>,
191-
options?: ZodFetchOptions
192-
): Promise<ZodFetchResult<z.output<TResponseBodySchema>>> {
211+
options?: ZodFetchOptions<z.output<TResponseBodySchema>, TOutput>
212+
): Promise<ZodFetchResult<TOutput>> {
193213
let $requestInit = await requestInit;
194214

195215
return traceZodFetch({ url, requestInit: $requestInit, options }, async (span) => {
@@ -202,7 +222,7 @@ async function _doZodFetch<TResponseBodySchema extends z.ZodTypeAny>(
202222
}
203223

204224
if (options?.prepareData) {
205-
result.data = await options.prepareData(result.data);
225+
result.data = await options.prepareData(result.data, result.response);
206226
}
207227

208228
return result;
@@ -707,7 +727,7 @@ export async function wrapZodFetch<T extends z.ZodTypeAny>(
707727
schema: T,
708728
url: string,
709729
requestInit?: RequestInit,
710-
options?: ZodFetchOptions<z.output<T>>
730+
options?: ZodFetchOptions<z.output<T>, z.infer<T>>
711731
): Promise<ApiResult<z.infer<T>>> {
712732
try {
713733
const response = await zodfetch(schema, url, requestInit, {

‎packages/core/src/v3/apiClient/index.ts

Lines changed: 79 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -214,36 +214,37 @@ export class ApiClient {
214214
headers: this.#getHeaders(clientOptions?.spanParentAsLink ?? false),
215215
body: JSON.stringify(body),
216216
},
217-
mergeRequestOptions(this.defaultRequestOptions, requestOptions)
218-
)
219-
.withResponse()
220-
.then(async ({ response, data }) => {
221-
const jwtHeader = response.headers.get("x-trigger-jwt");
217+
{
218+
...mergeRequestOptions(this.defaultRequestOptions, requestOptions),
219+
prepareData: async (data, response) => {
220+
const jwtHeader = response.headers.get("x-trigger-jwt");
221+
222+
if (typeof jwtHeader === "string") {
223+
return {
224+
...data,
225+
publicAccessToken: jwtHeader,
226+
};
227+
}
228+
229+
const claimsHeader = response.headers.get("x-trigger-jwt-claims");
230+
const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
231+
232+
const jwt = await generateJWT({
233+
secretKey: this.accessToken,
234+
payload: {
235+
...claims,
236+
scopes: [`read:runs:${data.id}`],
237+
},
238+
expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
239+
});
222240

223-
if (typeof jwtHeader === "string") {
224241
return {
225242
...data,
226-
publicAccessToken: jwtHeader,
243+
publicAccessToken: jwt,
227244
};
228-
}
229-
230-
const claimsHeader = response.headers.get("x-trigger-jwt-claims");
231-
const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
232-
233-
const jwt = await generateJWT({
234-
secretKey: this.accessToken,
235-
payload: {
236-
...claims,
237-
scopes: [`read:runs:${data.id}`],
238-
},
239-
expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
240-
});
241-
242-
return {
243-
...data,
244-
publicAccessToken: jwt,
245-
};
246-
});
245+
},
246+
}
247+
);
247248
}
248249

249250
batchTriggerV3(
@@ -261,27 +262,28 @@ export class ApiClient {
261262
}),
262263
body: JSON.stringify(body),
263264
},
264-
mergeRequestOptions(this.defaultRequestOptions, requestOptions)
265-
)
266-
.withResponse()
267-
.then(async ({ response, data }) => {
268-
const claimsHeader = response.headers.get("x-trigger-jwt-claims");
269-
const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
270-
271-
const jwt = await generateJWT({
272-
secretKey: this.accessToken,
273-
payload: {
274-
...claims,
275-
scopes: [`read:batch:${data.id}`],
276-
},
277-
expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
278-
});
279-
280-
return {
281-
...data,
282-
publicAccessToken: jwt,
283-
};
284-
});
265+
{
266+
...mergeRequestOptions(this.defaultRequestOptions, requestOptions),
267+
prepareData: async (data, response) => {
268+
const claimsHeader = response.headers.get("x-trigger-jwt-claims");
269+
const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
270+
271+
const jwt = await generateJWT({
272+
secretKey: this.accessToken,
273+
payload: {
274+
...claims,
275+
scopes: [`read:batch:${data.id}`],
276+
},
277+
expirationTime: requestOptions?.publicAccessToken?.expirationTime ?? "1h",
278+
});
279+
280+
return {
281+
...data,
282+
publicAccessToken: jwt,
283+
};
284+
},
285+
}
286+
);
285287
}
286288

287289
createUploadPayloadUrl(filename: string, requestOptions?: ZodFetchOptions) {
@@ -669,7 +671,36 @@ export class ApiClient {
669671
headers: this.#getHeaders(false),
670672
body: JSON.stringify(options),
671673
},
672-
mergeRequestOptions(this.defaultRequestOptions, requestOptions)
674+
{
675+
...mergeRequestOptions(this.defaultRequestOptions, requestOptions),
676+
prepareData: async (data, response) => {
677+
const jwtHeader = response.headers.get("x-trigger-jwt");
678+
679+
if (typeof jwtHeader === "string") {
680+
return {
681+
...data,
682+
publicAccessToken: jwtHeader,
683+
};
684+
}
685+
686+
const claimsHeader = response.headers.get("x-trigger-jwt-claims");
687+
const claims = claimsHeader ? JSON.parse(claimsHeader) : undefined;
688+
689+
const jwt = await generateJWT({
690+
secretKey: this.accessToken,
691+
payload: {
692+
...claims,
693+
scopes: [`write:waitpoints:${data.id}`],
694+
},
695+
expirationTime: "1h",
696+
});
697+
698+
return {
699+
...data,
700+
publicAccessToken: jwt,
701+
};
702+
},
703+
}
673704
);
674705
}
675706

‎packages/trigger-sdk/src/v3/wait.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,10 +48,14 @@ import { SpanStatusCode } from "@opentelemetry/api";
4848
* @param requestOptions - The request options for the waitpoint token.
4949
* @returns The waitpoint token.
5050
*/
51+
export type CreateWaitpointTokenResponse = CreateWaitpointTokenResponseBody & {
52+
publicAccessToken: string;
53+
};
54+
5155
function createToken(
5256
options?: CreateWaitpointTokenRequestBody,
5357
requestOptions?: ApiRequestOptions
54-
): ApiPromise<CreateWaitpointTokenResponseBody> {
58+
): ApiPromise<CreateWaitpointTokenResponse> {
5559
const apiClient = apiClientManager.clientOrThrow();
5660

5761
const $requestOptions = mergeRequestOptions(

0 commit comments

Comments
 (0)