feat: improve memory management

demo/go.mod

@@ -11,8 +11,8 @@ require (
 	github.com/rs/cors v1.11.0
 	github.com/vektah/gqlparser/v2 v2.5.30
 	github.com/wundergraph/cosmo/composition-go v0.0.0-20250820135159-bf8852195d3f
-	github.com/wundergraph/cosmo/router v0.0.0-20251125205644-175f80c4e6d9
-	github.com/wundergraph/cosmo/router-tests v0.0.0-20251125205644-175f80c4e6d9
+	github.com/wundergraph/cosmo/router v0.0.0-20251126113850-ce5751591be9
+	github.com/wundergraph/cosmo/router-tests v0.0.0-20251126113850-ce5751591be9
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.58.0
 	go.opentelemetry.io/otel v1.36.0
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.33.0
@@ -146,9 +146,10 @@ require (
 	github.com/urfave/cli/v2 v2.27.7 // indirect
 	github.com/vbatts/tar-split v0.12.1 // indirect
 	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
-	github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083 // indirect
+	github.com/wundergraph/astjson v1.0.0 // indirect
 	github.com/wundergraph/cosmo/demo/pkg/subgraphs/projects v0.0.0-20250715110703-10f2e5f9c79e // indirect
-	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.239 // indirect
+	github.com/wundergraph/go-arena v1.1.0 // indirect
+	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.239.0.20251126110525-c8b5567a9dd1 // indirect
 	github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	github.com/yusufpapurcu/wmi v1.2.4 // indirect

demo/go.sum

@@ -373,18 +373,20 @@ github.com/vektah/gqlparser/v2 v2.5.30 h1:EqLwGAFLIzt1wpx1IPpY67DwUujF1OfzgEyDsL
 github.com/vektah/gqlparser/v2 v2.5.30/go.mod h1:D1/VCZtV3LPnQrcPBeR/q5jkSQIPti0uYCP/RI0gIeo=
 github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
 github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
-github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083 h1:8/D7f8gKxTBjW+SZK4mhxTTBVpxcqeBgWF1Rfmltbfk=
-github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083/go.mod h1:eOTL6acwctsN4F3b7YE+eE2t8zcJ/doLm9sZzsxxxrE=
+github.com/wundergraph/astjson v1.0.0 h1:rETLJuQkMWWW03HCF6WBttEBOu8gi5vznj5KEUPVV2Q=
+github.com/wundergraph/astjson v1.0.0/go.mod h1:h12D/dxxnedtLzsKyBLK7/Oe4TAoGpRVC9nDpDrZSWw=
 github.com/wundergraph/cosmo/composition-go v0.0.0-20250820135159-bf8852195d3f h1:TjxywtqT+D6Qa/WGg1Vqn4ogT6P7M+iuUxPgDds1/kU=
 github.com/wundergraph/cosmo/composition-go v0.0.0-20250820135159-bf8852195d3f/go.mod h1:S3y/4/28Iut0SHMbmfXXa/6qyP55BA2rS4K3VI8A4Pg=
 github.com/wundergraph/cosmo/demo/pkg/subgraphs/projects v0.0.0-20250715110703-10f2e5f9c79e h1:VdJNlsiyWYxJzAD3jEe+DAQdzxkf9btD8qQNYNU+xQU=
 github.com/wundergraph/cosmo/demo/pkg/subgraphs/projects v0.0.0-20250715110703-10f2e5f9c79e/go.mod h1:WZ0yBeaDSGHqDMcQrP1JRYgCj9atF7ORXF8srnd2Sro=
-github.com/wundergraph/cosmo/router v0.0.0-20251125205644-175f80c4e6d9 h1:3bbV5IG402Od27mUDrMoLpA3kXdmLxX15CdYU91WECY=
-github.com/wundergraph/cosmo/router v0.0.0-20251125205644-175f80c4e6d9/go.mod h1:Xvxq+kCu1SHtaMueH7IK+ARX6+Bf9uKl00CZCl/LoOI=
-github.com/wundergraph/cosmo/router-tests v0.0.0-20251125205644-175f80c4e6d9 h1:gdaogSkGxTtZkUjunQ3mHdvTEpPV+fswvrQEOxxKRPA=
-github.com/wundergraph/cosmo/router-tests v0.0.0-20251125205644-175f80c4e6d9/go.mod h1:KVSvF1ivyKqmy+3TGl8QoI+sUWwZswcqXjZtApszpzA=
-github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.239 h1:wh8qTtVS4Wr/dJ/s162hAvCPsaZ1VOdmmg82QhNcGBE=
-github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.239/go.mod h1:mX25ASEQiKamxaFSK6NZihh0oDCigIuzro30up4mFH4=
+github.com/wundergraph/cosmo/router v0.0.0-20251126113850-ce5751591be9 h1:nfdYCTux8XROt0IVSHcN9t+FtELoC+zhWmZJGw/+SBk=
+github.com/wundergraph/cosmo/router v0.0.0-20251126113850-ce5751591be9/go.mod h1:Xvxq+kCu1SHtaMueH7IK+ARX6+Bf9uKl00CZCl/LoOI=
+github.com/wundergraph/cosmo/router-tests v0.0.0-20251126113850-ce5751591be9 h1:yPQvqWC6fY1cFyc/fiXf5KwXR9xXjWyfu0QR76cQWx8=
+github.com/wundergraph/cosmo/router-tests v0.0.0-20251126113850-ce5751591be9/go.mod h1:w072sJl/2jt4ZgV6rH5iX5xa46o51CbBJF3PvjeCm8I=
+github.com/wundergraph/go-arena v1.1.0 h1:9+wSRkJAkA2vbYHp6s8tEGhPViRGQNGXqPHT0QzhdIc=
+github.com/wundergraph/go-arena v1.1.0/go.mod h1:ROOysEHWJjLQ8FSfNxZCziagb7Qw2nXY3/vgKRh7eWw=
+github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.239.0.20251126110525-c8b5567a9dd1 h1:ILzek5T1ImcrFWs02kWrJ+fFhGnhFdP4yFWPyU6FIzo=
+github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.239.0.20251126110525-c8b5567a9dd1/go.mod h1:MFbY0QI8ncF60DHs7yyyiyyhWyld0WE1JokiyTVY8j4=
 github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342 h1:FnBeRrxr7OU4VvAzt5X7s6266i6cSVkkFPS0TuXWbIg=
 github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=

demo/pkg/subgraphs/courses/generated/mapping.json

@@ -53,22 +53,6 @@
     }
   ],
   "entityMappings": [
-    {
-      "typeName": "Course",
-      "kind": "entity",
-      "key": "id",
-      "rpc": "LookupCourseById",
-      "request": "LookupCourseByIdRequest",
-      "response": "LookupCourseByIdResponse"
-    },
-    {
-      "typeName": "Lesson",
-      "kind": "entity",
-      "key": "id",
-      "rpc": "LookupLessonById",
-      "request": "LookupLessonByIdRequest",
-      "response": "LookupLessonByIdResponse"
-    },
     {
       "typeName": "Employee",
       "kind": "entity",

demo/pkg/subgraphs/courses/generated/service.proto

@@ -5,12 +5,8 @@ import "google/protobuf/wrappers.proto";
 
 // Service definition for CoursesService
 service CoursesService {
-  // Lookup Course entity by id
-  rpc LookupCourseById(LookupCourseByIdRequest) returns (LookupCourseByIdResponse) {}
   // Lookup Employee entity by id
   rpc LookupEmployeeById(LookupEmployeeByIdRequest) returns (LookupEmployeeByIdResponse) {}
-  // Lookup Lesson entity by id
-  rpc LookupLessonById(LookupLessonByIdRequest) returns (LookupLessonByIdResponse) {}
   rpc MutationAddCourse(MutationAddCourseRequest) returns (MutationAddCourseResponse) {}
   rpc MutationAddLesson(MutationAddLessonRequest) returns (MutationAddLessonResponse) {}
   rpc QueryCourse(QueryCourseRequest) returns (QueryCourseResponse) {}
@@ -20,74 +16,6 @@ service CoursesService {
   rpc QueryThrowErrorCourses(QueryThrowErrorCoursesRequest) returns (QueryThrowErrorCoursesResponse) {}
 }
 
-// Key message for Course entity lookup
-message LookupCourseByIdRequestKey {
-  // Key field for Course entity lookup.
-  string id = 1;
-}
-
-// Request message for Course entity lookup.
-message LookupCourseByIdRequest {
-  /*
-   * List of keys to look up Course entities.
-   * Order matters - each key maps to one entity in LookupCourseByIdResponse.
-   */
-  repeated LookupCourseByIdRequestKey keys = 1;
-}
-
-// Response message for Course entity lookup.
-message LookupCourseByIdResponse {
-  /*
-   * List of Course entities in the same order as the keys in LookupCourseByIdRequest.
-   * Always return the same number of entities as keys. Use null for entities that cannot be found.
-   * 
-   * Example:
-   *   LookupUserByIdRequest:
-   *     keys:
-   *       - id: 1
-   *       - id: 2
-   *   LookupUserByIdResponse:
-   *     result:
-   *       - id: 1 # User with id 1 found
-   *       - null  # User with id 2 not found
-   */
-  repeated Course result = 1;
-}
-
-// Key message for Lesson entity lookup
-message LookupLessonByIdRequestKey {
-  // Key field for Lesson entity lookup.
-  string id = 1;
-}
-
-// Request message for Lesson entity lookup.
-message LookupLessonByIdRequest {
-  /*
-   * List of keys to look up Lesson entities.
-   * Order matters - each key maps to one entity in LookupLessonByIdResponse.
-   */
-  repeated LookupLessonByIdRequestKey keys = 1;
-}
-
-// Response message for Lesson entity lookup.
-message LookupLessonByIdResponse {
-  /*
-   * List of Lesson entities in the same order as the keys in LookupLessonByIdRequest.
-   * Always return the same number of entities as keys. Use null for entities that cannot be found.
-   * 
-   * Example:
-   *   LookupUserByIdRequest:
-   *     keys:
-   *       - id: 1
-   *       - id: 2
-   *   LookupUserByIdResponse:
-   *     result:
-   *       - id: 1 # User with id 1 found
-   *       - null  # User with id 2 not found
-   */
-  repeated Lesson result = 1;
-}
-
 // Key message for Employee entity lookup
 message LookupEmployeeByIdRequestKey {
   // Key field for Employee entity lookup.
@@ -179,6 +107,11 @@ message MutationAddLessonResponse {
   Lesson add_lesson = 1;
 }
 
+message Employee {
+  int32 id = 1;
+  repeated Course taught_courses = 2;
+}
+
 message Course {
   string id = 1;
   string title = 2;
@@ -194,9 +127,4 @@ message Lesson {
   google.protobuf.StringValue description = 4;
   int32 order = 5;
   Course course = 6;
-}
-
-message Employee {
-  int32 id = 1;
-  repeated Course taught_courses = 2;
 }

demo/pkg/subgraphs/courses/generated/service_grpc_pb.d.ts

@@ -9,9 +9,7 @@ import * as service_pb from "./service_pb";
 import * as google_protobuf_wrappers_pb from "google-protobuf/google/protobuf/wrappers_pb";
 
 interface ICoursesServiceService extends grpc.ServiceDefinition<grpc.UntypedServiceImplementation> {
-    lookupCourseById: ICoursesServiceService_ILookupCourseById;
     lookupEmployeeById: ICoursesServiceService_ILookupEmployeeById;
-    lookupLessonById: ICoursesServiceService_ILookupLessonById;
     mutationAddCourse: ICoursesServiceService_IMutationAddCourse;
     mutationAddLesson: ICoursesServiceService_IMutationAddLesson;
     queryCourse: ICoursesServiceService_IQueryCourse;
@@ -21,15 +19,6 @@ interface ICoursesServiceService extends grpc.ServiceDefinition<grpc.UntypedServ
     queryThrowErrorCourses: ICoursesServiceService_IQueryThrowErrorCourses;
 }
 
-interface ICoursesServiceService_ILookupCourseById extends grpc.MethodDefinition<service_pb.LookupCourseByIdRequest, service_pb.LookupCourseByIdResponse> {
-    path: "/service.CoursesService/LookupCourseById";
-    requestStream: false;
-    responseStream: false;
-    requestSerialize: grpc.serialize<service_pb.LookupCourseByIdRequest>;
-    requestDeserialize: grpc.deserialize<service_pb.LookupCourseByIdRequest>;
-    responseSerialize: grpc.serialize<service_pb.LookupCourseByIdResponse>;
-    responseDeserialize: grpc.deserialize<service_pb.LookupCourseByIdResponse>;
-}
 interface ICoursesServiceService_ILookupEmployeeById extends grpc.MethodDefinition<service_pb.LookupEmployeeByIdRequest, service_pb.LookupEmployeeByIdResponse> {
     path: "/service.CoursesService/LookupEmployeeById";
     requestStream: false;
@@ -39,15 +28,6 @@ interface ICoursesServiceService_ILookupEmployeeById extends grpc.MethodDefiniti
     responseSerialize: grpc.serialize<service_pb.LookupEmployeeByIdResponse>;
     responseDeserialize: grpc.deserialize<service_pb.LookupEmployeeByIdResponse>;
 }
-interface ICoursesServiceService_ILookupLessonById extends grpc.MethodDefinition<service_pb.LookupLessonByIdRequest, service_pb.LookupLessonByIdResponse> {
-    path: "/service.CoursesService/LookupLessonById";
-    requestStream: false;
-    responseStream: false;
-    requestSerialize: grpc.serialize<service_pb.LookupLessonByIdRequest>;
-    requestDeserialize: grpc.deserialize<service_pb.LookupLessonByIdRequest>;
-    responseSerialize: grpc.serialize<service_pb.LookupLessonByIdResponse>;
-    responseDeserialize: grpc.deserialize<service_pb.LookupLessonByIdResponse>;
-}
 interface ICoursesServiceService_IMutationAddCourse extends grpc.MethodDefinition<service_pb.MutationAddCourseRequest, service_pb.MutationAddCourseResponse> {
     path: "/service.CoursesService/MutationAddCourse";
     requestStream: false;
@@ -115,9 +95,7 @@ interface ICoursesServiceService_IQueryThrowErrorCourses extends grpc.MethodDefi
 export const CoursesServiceService: ICoursesServiceService;
 
 export interface ICoursesServiceServer extends grpc.UntypedServiceImplementation {
-    lookupCourseById: grpc.handleUnaryCall<service_pb.LookupCourseByIdRequest, service_pb.LookupCourseByIdResponse>;
     lookupEmployeeById: grpc.handleUnaryCall<service_pb.LookupEmployeeByIdRequest, service_pb.LookupEmployeeByIdResponse>;
-    lookupLessonById: grpc.handleUnaryCall<service_pb.LookupLessonByIdRequest, service_pb.LookupLessonByIdResponse>;
     mutationAddCourse: grpc.handleUnaryCall<service_pb.MutationAddCourseRequest, service_pb.MutationAddCourseResponse>;
     mutationAddLesson: grpc.handleUnaryCall<service_pb.MutationAddLessonRequest, service_pb.MutationAddLessonResponse>;
     queryCourse: grpc.handleUnaryCall<service_pb.QueryCourseRequest, service_pb.QueryCourseResponse>;
@@ -128,15 +106,9 @@ export interface ICoursesServiceServer extends grpc.UntypedServiceImplementation
 }
 
 export interface ICoursesServiceClient {
-    lookupCourseById(request: service_pb.LookupCourseByIdRequest, callback: (error: grpc.ServiceError | null, response: service_pb.LookupCourseByIdResponse) => void): grpc.ClientUnaryCall;
-    lookupCourseById(request: service_pb.LookupCourseByIdRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.LookupCourseByIdResponse) => void): grpc.ClientUnaryCall;
-    lookupCourseById(request: service_pb.LookupCourseByIdRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.LookupCourseByIdResponse) => void): grpc.ClientUnaryCall;
     lookupEmployeeById(request: service_pb.LookupEmployeeByIdRequest, callback: (error: grpc.ServiceError | null, response: service_pb.LookupEmployeeByIdResponse) => void): grpc.ClientUnaryCall;
     lookupEmployeeById(request: service_pb.LookupEmployeeByIdRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.LookupEmployeeByIdResponse) => void): grpc.ClientUnaryCall;
     lookupEmployeeById(request: service_pb.LookupEmployeeByIdRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.LookupEmployeeByIdResponse) => void): grpc.ClientUnaryCall;
-    lookupLessonById(request: service_pb.LookupLessonByIdRequest, callback: (error: grpc.ServiceError | null, response: service_pb.LookupLessonByIdResponse) => void): grpc.ClientUnaryCall;
-    lookupLessonById(request: service_pb.LookupLessonByIdRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.LookupLessonByIdResponse) => void): grpc.ClientUnaryCall;
-    lookupLessonById(request: service_pb.LookupLessonByIdRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.LookupLessonByIdResponse) => void): grpc.ClientUnaryCall;
     mutationAddCourse(request: service_pb.MutationAddCourseRequest, callback: (error: grpc.ServiceError | null, response: service_pb.MutationAddCourseResponse) => void): grpc.ClientUnaryCall;
     mutationAddCourse(request: service_pb.MutationAddCourseRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.MutationAddCourseResponse) => void): grpc.ClientUnaryCall;
     mutationAddCourse(request: service_pb.MutationAddCourseRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.MutationAddCourseResponse) => void): grpc.ClientUnaryCall;
@@ -162,15 +134,9 @@ export interface ICoursesServiceClient {
 
 export class CoursesServiceClient extends grpc.Client implements ICoursesServiceClient {
     constructor(address: string, credentials: grpc.ChannelCredentials, options?: Partial<grpc.ClientOptions>);
-    public lookupCourseById(request: service_pb.LookupCourseByIdRequest, callback: (error: grpc.ServiceError | null, response: service_pb.LookupCourseByIdResponse) => void): grpc.ClientUnaryCall;
-    public lookupCourseById(request: service_pb.LookupCourseByIdRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.LookupCourseByIdResponse) => void): grpc.ClientUnaryCall;
-    public lookupCourseById(request: service_pb.LookupCourseByIdRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.LookupCourseByIdResponse) => void): grpc.ClientUnaryCall;
     public lookupEmployeeById(request: service_pb.LookupEmployeeByIdRequest, callback: (error: grpc.ServiceError | null, response: service_pb.LookupEmployeeByIdResponse) => void): grpc.ClientUnaryCall;
     public lookupEmployeeById(request: service_pb.LookupEmployeeByIdRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.LookupEmployeeByIdResponse) => void): grpc.ClientUnaryCall;
     public lookupEmployeeById(request: service_pb.LookupEmployeeByIdRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.LookupEmployeeByIdResponse) => void): grpc.ClientUnaryCall;
-    public lookupLessonById(request: service_pb.LookupLessonByIdRequest, callback: (error: grpc.ServiceError | null, response: service_pb.LookupLessonByIdResponse) => void): grpc.ClientUnaryCall;
-    public lookupLessonById(request: service_pb.LookupLessonByIdRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.LookupLessonByIdResponse) => void): grpc.ClientUnaryCall;
-    public lookupLessonById(request: service_pb.LookupLessonByIdRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.LookupLessonByIdResponse) => void): grpc.ClientUnaryCall;
     public mutationAddCourse(request: service_pb.MutationAddCourseRequest, callback: (error: grpc.ServiceError | null, response: service_pb.MutationAddCourseResponse) => void): grpc.ClientUnaryCall;
     public mutationAddCourse(request: service_pb.MutationAddCourseRequest, metadata: grpc.Metadata, callback: (error: grpc.ServiceError | null, response: service_pb.MutationAddCourseResponse) => void): grpc.ClientUnaryCall;
     public mutationAddCourse(request: service_pb.MutationAddCourseRequest, metadata: grpc.Metadata, options: Partial<grpc.CallOptions>, callback: (error: grpc.ServiceError | null, response: service_pb.MutationAddCourseResponse) => void): grpc.ClientUnaryCall;