Add support for async derived context functions

src/Errors.ts

@@ -255,8 +255,8 @@ export function methodMissingType() {
   return "Expected GraphQL field methods to have an explicitly defined return type. Grats needs to be able to see the type of the field to generate its type in the GraphQL schema.";
 }
 
-export function wrapperMissingTypeArg() {
-  return `Expected wrapper type reference to have type arguments. Grats needs to be able to see the return type in order to generate a GraphQL schema.`;
+export function wrapperMissingTypeArg(wrapperTypeName: string) {
+  return `Expected \`${wrapperTypeName}\` type to have exactly one type argument. Grats needs to be able to see the inner type in order to generate a GraphQL schema.`;
 }
 
 export function invalidWrapperOnInputType(wrapperName: string) {
@@ -607,7 +607,7 @@ export function invalidDerivedContextArgType() {
 }
 
 export function missingReturnTypeForDerivedResolver() {
-  return 'Expected derived resolver to have an explicit return type. This is needed to allow Grats to "see" which type to treat as a derived context type.';
+  return 'Expected derived resolver\'s return type to be a named type alias, e.g. `: SomeType`. This is needed to allow Grats to "see" which type declaration to treat as the derived context type.';
 }
 
 export function derivedResolverInvalidReturnType() {

src/Extractor.ts

@@ -411,8 +411,14 @@ class Extractor {
     if (returnType == null) {
       return this.report(node, E.missingReturnTypeForDerivedResolver());
     }
-    if (!ts.isTypeReferenceNode(returnType)) {
-      return this.report(returnType, E.missingReturnTypeForDerivedResolver());
+
+    // Check if the return type is Promise<T> and unwrap it
+    const unwrapped = this.maybeUnwrapPromiseType(returnType);
+    if (unwrapped === null) return null;
+    const { type: innerType, isAsync } = unwrapped;
+
+    if (!ts.isTypeReferenceNode(innerType)) {
+      return this.report(innerType, E.missingReturnTypeForDerivedResolver());
     }
 
     const funcName = this.namedFunctionExportName(node);
@@ -434,8 +440,9 @@ class Extractor {
         path: tsModulePath,
         exportName: funcName?.text ?? null,
         args: paramResults.resolverParams,
+        async: isAsync,
       },
-      returnType,
+      innerType,
     );
   }
 
@@ -2592,6 +2599,30 @@ class Extractor {
     return null;
   }
 
+  /**
+   * Unwraps a Promise<T> type to T, tracking whether it was async.
+   * Returns null if there's an error (e.g., Promise without type arguments).
+   */
+  maybeUnwrapPromiseType(
+    type: ts.TypeNode,
+  ): { type: ts.TypeNode; isAsync: boolean } | null {
+    if (!ts.isTypeReferenceNode(type)) {
+      return { type, isAsync: false };
+    }
+
+    const typeName = type.typeName;
+    if (ts.isIdentifier(typeName) && typeName.text === "Promise") {
+      if (type.typeArguments == null || type.typeArguments.length !== 1) {
+        //
+        this.report(type, E.wrapperMissingTypeArg(typeName.text));
+        return null;
+      }
+      return { type: type.typeArguments[0], isAsync: true };
+    }
+
+    return { type, isAsync: false };
+  }
+
   typeReference(
     node: ts.TypeReferenceNode,
     ctx: FieldTypeContext,
@@ -2632,10 +2663,9 @@ class Extractor {
         return this.gql.nonNullType(node, listType);
       }
       case "Promise": {
-        if (node.typeArguments == null) {
-          return this.report(node, E.wrapperMissingTypeArg());
-        }
-        const element = this.collectType(node.typeArguments[0], ctx);
+        const unwrapped = this.maybeUnwrapPromiseType(node);
+        if (unwrapped === null) return null;
+        const element = this.collectType(unwrapped.type, ctx);
         if (element == null) return null;
         return element;
       }

src/TypeContext.ts

@@ -27,6 +27,7 @@ export type DerivedResolverDefinition = {
   exportName: string | null;
   args: ResolverArgument[];
   kind: "DERIVED_CONTEXT";
+  async: boolean;
 };
 
 export type NameDefinition = {

src/codegen/TSAstBuilder.ts

@@ -56,9 +56,10 @@ export default class TSAstBuilder {
     name: string,
     params: ts.ParameterDeclaration[],
     statements: ts.Statement[],
+    isAsync: boolean = false,
   ): ts.MethodDeclaration {
     return F.createMethodDeclaration(
-      undefined,
+      isAsync ? [F.createModifier(ts.SyntaxKind.AsyncKeyword)] : undefined,
       undefined,
       name,
       undefined,

src/codegen/resolverCodegen.ts

@@ -64,11 +64,14 @@ export default class ResolverCodegen {
               ),
             ),
           ],
+          false,
         );
       case "method": {
+        const args = resolver.arguments ?? [];
+        const isAsync = args.some(usesAsyncDerivedContext);
         return this.ts.method(
           methodName,
-          extractUsedParams(resolver.arguments ?? [], true).map((name) => {
+          extractUsedParams(args, true).map((name) => {
             if (name === "source") {
               return this.ts.param("source", getSourceTypeRef());
             }
@@ -82,12 +85,13 @@ export default class ResolverCodegen {
                   F.createIdentifier(resolver.name ?? fieldName),
                 ),
                 [],
-                (resolver.arguments ?? []).map((arg) => {
+                args.map((arg) => {
                   return this.resolverParam(arg);
                 }),
               ),
             ),
           ],
+          isAsync,
         );
       }
       case "function": {
@@ -101,9 +105,11 @@ export default class ResolverCodegen {
           resolverName,
           false,
         );
+        const args = resolver.arguments ?? [];
+        const isAsync = args.some(usesAsyncDerivedContext);
         return this.ts.method(
           methodName,
-          extractUsedParams(resolver.arguments ?? []).map((name) => {
+          extractUsedParams(args).map((name) => {
             if (name === "source") {
               return this.ts.param("source", getSourceTypeRef());
             }
@@ -114,12 +120,13 @@ export default class ResolverCodegen {
               F.createCallExpression(
                 F.createIdentifier(resolverName),
                 undefined,
-                (resolver.arguments ?? []).map((arg) => {
+                args.map((arg) => {
                   return this.resolverParam(arg);
                 }),
               ),
             ),
           ],
+          isAsync,
         );
       }
       case "staticMethod": {
@@ -135,9 +142,11 @@ export default class ResolverCodegen {
           resolverName,
           false,
         );
+        const args = resolver.arguments ?? [];
+        const isAsync = args.some(usesAsyncDerivedContext);
         return this.ts.method(
           methodName,
-          extractUsedParams(resolver.arguments ?? []).map((name) => {
+          extractUsedParams(args).map((name) => {
             if (name === "source") {
               return this.ts.param("source", getSourceTypeRef());
             }
@@ -151,12 +160,13 @@ export default class ResolverCodegen {
                   F.createIdentifier(resolver.name),
                 ),
                 undefined,
-                (resolver.arguments ?? []).map((arg) => {
+                args.map((arg) => {
                   return this.resolverParam(arg);
                 }),
               ),
             ),
           ],
+          isAsync,
         );
       }
       default:
@@ -213,11 +223,16 @@ export default class ResolverCodegen {
       case "derivedContext": {
         const localName = this.getDerivedContextName(arg.path, arg.exportName);
         this.ts.importUserConstruct(arg.path, arg.exportName, localName, false);
-        return F.createCallExpression(
+        const callExpr = F.createCallExpression(
           F.createIdentifier(localName),
           undefined,
           arg.args.map((arg) => this.resolverParam(arg)),
         );
+        // If the derived context is async, we need to await it
+        if (arg.async) {
+          return F.createAwaitExpression(callExpr);
+        }
+        return callExpr;
       }
 
       default:
@@ -361,6 +376,16 @@ function usesContext(param: ResolverArgument) {
   }
 }
 
+// Check if any param is or uses an async derived context
+function usesAsyncDerivedContext(param: ResolverArgument): boolean {
+  switch (param.kind) {
+    case "derivedContext":
+      return param.async || param.args.some(usesAsyncDerivedContext);
+    default:
+      return false;
+  }
+}
+
 function fieldDirective(
   field: GraphQLField<unknown, unknown>,
   name: string,

src/metadata.ts

@@ -114,6 +114,7 @@ export type DerivedContextArgument = {
   path: string; // Path to the module
   exportName: string | null; // Export name. If omitted, the class is the default export
   args: Array<ContextArgs>;
+  async: boolean;
 };
 
 /** The GraphQL info object */

src/resolverSignature.ts

@@ -66,6 +66,7 @@ export type DerivedContextResolverArgument = {
   exportName: string | null;
   args: Array<DerivedContextResolverArgument | ContextResolverArgument>;
   node: ts.Node;
+  async: boolean;
 };
 
 export type InformationResolverArgument = {

src/tests/fixtures/derived_context/derivedContextNoReturnType.invalid.ts.expected.md

@@ -29,7 +29,7 @@ export function greeting(_: Query, ctx: DerivedContext): string {
 ### Error Report
 
 ```text
-src/tests/fixtures/derived_context/derivedContextNoReturnType.invalid.ts:11:1 - error: Expected derived resolver to have an explicit return type. This is needed to allow Grats to "see" which type to treat as a derived context type.
+src/tests/fixtures/derived_context/derivedContextNoReturnType.invalid.ts:11:1 - error: Expected derived resolver's return type to be a named type alias, e.g. `: SomeType`. This is needed to allow Grats to "see" which type declaration to treat as the derived context type.
 
 11 export function createDerivedContext(ctx: RootContext) {
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

src/tests/fixtures/derived_context/derivedContextNonNamedReturnType.invalid.ts.expected.md

@@ -29,7 +29,7 @@ export function greeting(_: Query, ctx: DerivedContext): string {
 ### Error Report
 
 ```text
-src/tests/fixtures/derived_context/derivedContextNonNamedReturnType.invalid.ts:11:57 - error: Expected derived resolver to have an explicit return type. This is needed to allow Grats to "see" which type to treat as a derived context type.
+src/tests/fixtures/derived_context/derivedContextNonNamedReturnType.invalid.ts:11:57 - error: Expected derived resolver's return type to be a named type alias, e.g. `: SomeType`. This is needed to allow Grats to "see" which type declaration to treat as the derived context type.
 
 11 export function createDerivedContext(ctx: RootContext): { greeting: string } {
                                                            ~~~~~~~~~~~~~~~~~~~~

src/tests/fixtures/derived_context/derivedContextPromiseNoTypeArg.invalid.ts

@@ -0,0 +1,18 @@
+/** @gqlContext */
+type RootContext = {
+  userName: string;
+};
+
+type DerivedContext = {
+  greeting: string;
+};
+
+/** @gqlContext */
+export async function createDerivedContext(ctx: RootContext): Promise {
+  return { greeting: `Hello, ${ctx.userName}!` };
+}
+
+/** @gqlQueryField */
+export function greeting(ctx: DerivedContext): string {
+  return ctx.greeting;
+}

src/tests/fixtures/derived_context/derivedContextPromiseNoTypeArg.invalid.ts.expected.md

@@ -0,0 +1,33 @@
+## input
+
+```ts title="derived_context/derivedContextPromiseNoTypeArg.invalid.ts"
+/** @gqlContext */
+type RootContext = {
+  userName: string;
+};
+
+type DerivedContext = {
+  greeting: string;
+};
+
+/** @gqlContext */
+export async function createDerivedContext(ctx: RootContext): Promise {
+  return { greeting: `Hello, ${ctx.userName}!` };
+}
+
+/** @gqlQueryField */
+export function greeting(ctx: DerivedContext): string {
+  return ctx.greeting;
+}
+```
+
+## Output
+
+### Error Report
+
+```text
+src/tests/fixtures/derived_context/derivedContextPromiseNoTypeArg.invalid.ts:11:63 - error: Expected `Promise` type to have exactly one type argument. Grats needs to be able to see the inner type in order to generate a GraphQL schema.
+
+11 export async function createDerivedContext(ctx: RootContext): Promise {
+                                                                 ~~~~~~~
+```

src/tests/fixtures/derived_context/derivedContextPromiseNonReferenceInner.invalid.ts

@@ -0,0 +1,19 @@
+/** @gqlContext */
+type RootContext = {
+  userName: string;
+};
+
+/** @gqlContext */
+export async function createDerivedContext(
+  ctx: RootContext,
+): Promise<{ greeting: string }> {
+  return { greeting: `Hello, ${ctx.userName}!` };
+}
+
+/** @gqlType */
+type Query = unknown;
+
+/** @gqlField */
+export function greeting(_: Query, ctx: { greeting: string }): string {
+  return ctx.greeting;
+}

src/tests/fixtures/derived_context/derivedContextPromiseNonReferenceInner.invalid.ts.expected.md

@@ -0,0 +1,34 @@
+## input
+
+```ts title="derived_context/derivedContextPromiseNonReferenceInner.invalid.ts"
+/** @gqlContext */
+type RootContext = {
+  userName: string;
+};
+
+/** @gqlContext */
+export async function createDerivedContext(
+  ctx: RootContext,
+): Promise<{ greeting: string }> {
+  return { greeting: `Hello, ${ctx.userName}!` };
+}
+
+/** @gqlType */
+type Query = unknown;
+
+/** @gqlField */
+export function greeting(_: Query, ctx: { greeting: string }): string {
+  return ctx.greeting;
+}
+```
+
+## Output
+
+### Error Report
+
+```text
+src/tests/fixtures/derived_context/derivedContextPromiseNonReferenceInner.invalid.ts:9:12 - error: Expected derived resolver's return type to be a named type alias, e.g. `: SomeType`. This is needed to allow Grats to "see" which type declaration to treat as the derived context type.
+
+9 ): Promise<{ greeting: string }> {
+             ~~~~~~~~~~~~~~~~~~~~
+```