diff --git a/.changeset/fluffy-rivers-explode.md b/.changeset/fluffy-rivers-explode.md
new file mode 100644
index 0000000..c7bee30
--- /dev/null
+++ b/.changeset/fluffy-rivers-explode.md
@@ -0,0 +1,5 @@
+---
+"@labdigital/federated-token-yoga": minor
+---
+
+Require the FederatedToken to already exist on the context
diff --git a/.eslintrc.cjs b/.eslintrc.cjs
index b0cf6b4..f2f87af 100644
--- a/.eslintrc.cjs
+++ b/.eslintrc.cjs
@@ -12,7 +12,6 @@ module.exports = {
 			},
 		],
 		"unused-imports/no-unused-imports": "error",
-		"arrow-body-style": ["error", "as-needed"],
 		"no-console": [
 			"error",
 			{
diff --git a/packages/yoga/src/index.test.ts b/packages/yoga/src/index.test.ts
new file mode 100644
index 0000000..5008428
--- /dev/null
+++ b/packages/yoga/src/index.test.ts
@@ -0,0 +1,147 @@
+import { createSchema, createYoga } from "graphql-yoga";
+import { describe, it, expect } from "vitest";
+import { useFederatedToken } from ".";
+import { FederatedToken } from "@labdigital/federated-token";
+
+type FederatedTokenContext = {
+	federatedToken: FederatedToken;
+};
+
+export function assertSingleValue<TValue extends object>(
+	value: TValue | AsyncIterable<TValue>,
+): asserts value is TValue {
+	if (Symbol.asyncIterator in value) {
+		throw new Error("Expected single value");
+	}
+}
+
+describe("useFederatdToken()", () => {
+	const schemaFactory = async () => {
+		return createSchema<FederatedTokenContext>({
+			typeDefs: /* GraphQL */ `
+				type Query {
+					tokenData: String
+				}
+				type Mutation {
+					createToken(service: String!): Boolean
+				}
+			`,
+			resolvers: {
+				Query: {
+					tokenData: (parent, args, context) => {
+						return JSON.stringify({
+							tokens: context.federatedToken.tokens,
+							values: context.federatedToken.values,
+							refreshTokens: context.federatedToken.refreshTokens,
+						});
+					},
+				},
+				Mutation: {
+					createToken: (parent, args, context) => {
+						context.federatedToken.setAccessToken(args.service, {
+							token: "test",
+							exp: 123,
+							sub: "test",
+						});
+						context.federatedToken.setRefreshToken(
+							args.service,
+							"refresh-token",
+						);
+						return true;
+					},
+				},
+			},
+		});
+	};
+
+	const yoga = createYoga<FederatedTokenContext>({
+		schema: schemaFactory,
+		plugins: [useFederatedToken()],
+		context: () => {
+			return {
+				federatedToken: new FederatedToken(),
+			};
+		},
+	});
+
+	it("should handle missing tokens", async () => {
+		const response = await yoga.fetch("http://localhost:3000/graphql", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+			},
+			body: JSON.stringify({
+				query: "{ tokenData }",
+			}),
+		});
+		const content = await response.json();
+		assertSingleValue(content);
+	});
+
+	it("should load access and refresh tokens", async () => {
+		const token = new FederatedToken();
+		token.setValue("my-service", { foo: "bar" });
+		token.setAccessToken("my-service", { token: "bar", exp: 123, sub: "test" });
+		token.setRefreshToken("my-service", "refresh-token");
+
+		const response = await yoga.fetch("http://localhost:3000/graphql", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				"X-Access-Token": token.serializeAccessToken() ?? "",
+				"X-Refresh-Token": token.dumpRefreshToken() ?? "",
+			},
+			body: JSON.stringify({
+				query: "{ tokenData }",
+			}),
+		});
+		const content = await response.json();
+
+		assertSingleValue(content);
+
+		const result = JSON.parse(content.data.tokenData);
+		expect(result).toEqual({
+			tokens: {
+				"my-service": {
+					token: "bar",
+					exp: 123,
+					sub: "test",
+				},
+			},
+			values: {
+				"my-service": {
+					foo: "bar",
+				},
+			},
+			refreshTokens: {
+				"my-service": "refresh-token",
+			},
+		});
+
+		expect(response.headers.has("X-Access-Token")).toBe(false);
+		expect(response.headers.has("X-Refresh-Token")).toBe(false);
+	});
+
+	it("should create new token", async () => {
+		const response = await yoga.fetch("http://localhost:3000/graphql", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+			},
+			body: JSON.stringify({
+				query: `mutation { createToken(service: "foobar") }`,
+			}),
+		});
+		const content = await response.json();
+		assertSingleValue(content);
+
+		expect(response.headers.has("X-Access-Token")).toBe(true);
+		expect(response.headers.has("X-Refresh-Token")).toBe(true);
+
+		const accessToken = response.headers.get("X-Access-Token");
+		const refreshToken = response.headers.get("X-Refresh-Token");
+
+		expect(accessToken).toBeDefined();
+		expect(refreshToken).toBeDefined();
+	});
+});
diff --git a/packages/yoga/src/index.ts b/packages/yoga/src/index.ts
index 641c155..812260e 100644
--- a/packages/yoga/src/index.ts
+++ b/packages/yoga/src/index.ts
@@ -6,26 +6,26 @@ type FederatedTokenContext = {
 };
 
 export const useFederatedToken = <T extends FederatedTokenContext>(): Plugin<
-	object,
+	T,
 	T
 > => ({
-	onRequest: ({ request, serverContext }) => {
-		// Initialize FederatedToken and add it to the context
-		const federatedToken = serverContext.federatedToken ?? new FederatedToken();
-
-		// Retrieve tokens from headers using the serverContext
-		const accessToken = request.headers.get("x-access-token") as string;
-		const refreshToken = request.headers.get("x-refresh-token") as string;
+	onContextBuilding: ({ context, extendContext }) => {
+		if (!context.federatedToken) {
+			throw new Error("Federated token not found in context");
+		}
+		if (!context.request) {
+			throw new Error("Request not found in context");
+		}
 
+		const { request, federatedToken } = context;
+		const accessToken = request.headers.get("x-access-token");
+		const refreshToken = request.headers.get("x-refresh-token");
 		if (accessToken) {
 			federatedToken.deserializeAccessToken(accessToken);
 		}
-
 		if (refreshToken) {
 			federatedToken.loadRefreshToken(refreshToken);
 		}
-
-		serverContext.federatedToken = federatedToken;
 	},
 
 	onResponse: async ({ response, serverContext }) => {