fix: Packages, Typing, and Env Oh My

Author: Maximo-Guk

apps/cloudflare-one-casb/src/context.ts

@@ -0,0 +1,15 @@
+import type { CASBMCP, UserDetails } from './index'
+
+export interface Env {
+	ENVIRONMENT: 'development' | 'staging' | 'production'
+	MCP_SERVER_NAME: string
+	MCP_SERVER_VERSION: string
+	MCP_OBJECT: DurableObjectNamespace<CASBMCP>
+	MCP_METRICS: AnalyticsEngineDataset
+	AI: Ai
+
+	CLOUDFLARE_CLIENT_ID: string
+	CLOUDFLARE_CLIENT_SECRET: string
+
+	USER_DETAILS: DurableObjectNamespace<UserDetails>
+}

apps/cloudflare-one-casb/src/index.ts

@@ -1,19 +1,24 @@
 import OAuthProvider from '@cloudflare/workers-oauth-provider'
 import { McpAgent } from 'agents/mcp'
-import { env } from 'cloudflare:workers'
 
 import {
 	createAuthHandlers,
 	handleTokenExchangeCallback,
 } from '@repo/mcp-common/src/cloudflare-oauth-handler'
+import { getUserDetails, UserDetails } from '@repo/mcp-common/src/durable-objects/user_details'
+import { getEnv } from '@repo/mcp-common/src/env'
 import { CloudflareMCPServer } from '@repo/mcp-common/src/server'
 import { registerAccountTools } from '@repo/mcp-common/src/tools/account'
 
 import { MetricsTracker } from '../../../packages/mcp-observability/src'
 import { registerIntegrationsTools } from './tools/integrations'
 
 import type { AccountSchema, UserSchema } from '@repo/mcp-common/src/cloudflare-oauth-handler'
-import type { CloudflareMcpAgent } from '@repo/mcp-common/src/types/cloudflare-mcp-agent'
+import type { Env } from './context'
+
+export { UserDetails }
+
+const env = getEnv<Env>()
 
 const metrics = new MetricsTracker(env.MCP_METRICS, {
 	name: env.MCP_SERVER_NAME,
@@ -46,32 +51,37 @@ export class CASBMCP extends McpAgent<Env, State, Props> {
 	}
 
 	async init() {
-		this.server = new CloudflareMCPServer(this.props.user.id, this.env.MCP_METRICS, {
-			name: this.env.MCP_SERVER_NAME,
-			version: this.env.MCP_SERVER_VERSION,
+		this.server = new CloudflareMCPServer({
+			userId: this.props.user.id,
+			wae: this.env.MCP_METRICS,
+			serverInfo: {
+				name: this.env.MCP_SERVER_NAME,
+				version: this.env.MCP_SERVER_VERSION,
+			},
 		})
 
 		registerAccountTools(this)
 		registerIntegrationsTools(this)
 	}
 
-	getActiveAccountId() {
+	async getActiveAccountId() {
 		try {
-			return this.state.activeAccountId ?? null
+			// Get UserDetails Durable Object based off the userId and retrieve the activeAccountId from it
+			// we do this so we can persist activeAccountId across sessions
+			const userDetails = getUserDetails(env, this.props.user.id)
+			return await userDetails.getActiveAccountId()
 		} catch (e) {
-			console.error('getActiveAccountId failured: ', e)
+			this.server.recordError(e)
 			return null
 		}
 	}
 
-	setActiveAccountId(accountId: string) {
+	async setActiveAccountId(accountId: string) {
 		try {
-			this.setState({
-				...this.state,
-				activeAccountId: accountId,
-			})
+			const userDetails = getUserDetails(env, this.props.user.id)
+			await userDetails.setActiveAccountId(accountId)
 		} catch (e) {
-			return null
+			this.server.recordError(e)
 		}
 	}
 }

apps/cloudflare-one-casb/src/tools/integrations.ts

@@ -16,7 +16,7 @@ import {
 	assetCategoryVendorParam,
 } from '@repo/mcp-common/src/schemas/cf1-integrations'
 
-import type { McpAgentWithAccount } from '@repo/mcp-common/src/api/account'
+import type { CloudflareMcpAgent } from '@repo/mcp-common/src/types/cloudflare-mcp-agent'
 import type { ToolDefinition } from '@repo/mcp-common/src/types/tools'
 import type { CASBMCP } from '../index'
 
@@ -248,7 +248,7 @@ export function registerIntegrationsTools(agent: CASBMCP) {
 			name,
 			description,
 			params,
-			withAccountCheck(agent as unknown as McpAgentWithAccount, handler)
+			withAccountCheck(agent as unknown as CloudflareMcpAgent, handler)
 		)
 	})
 }

apps/cloudflare-one-casb/worker-configuration.d.ts

@@ -1,15 +1,14 @@
-// Generated by Wrangler by running `wrangler types` (hash: 085ce38c6421603ed1a73fa9678ae0a7)
+// Generated by Wrangler by running `wrangler types` (hash: 16591b3e7a816201c4869a97f0b7602f)
 // Runtime types generated with [email protected] 2025-03-10 nodejs_compat
 declare namespace Cloudflare {
 	interface Env {
 		OAUTH_KV: KVNamespace;
-		CLOUDFLARE_CLIENT_ID: string;
-		CLOUDFLARE_CLIENT_SECRET: string;
-		MCP_OBJECT: DurableObjectNamespace<import("./src/index").CASBMCP>;
-
-        MCP_SERVER_NAME: "PLACEHOLDER";
+		MCP_SERVER_NAME: "PLACEHOLDER";
 		MCP_SERVER_VERSION: "PLACEHOLDER";
-        MCP_METRICS: AnalyticsEngineDataset;
+		CLOUDFLARE_CLIENT_ID: "<PLACEHOLDER>";
+		CLOUDFLARE_CLIENT_SECRET: "<PLACEHOLDER>";
+		MCP_OBJECT: DurableObjectNamespace<import("./src/index").CASBMCP>;
+		MCP_METRICS: AnalyticsEngineDataset;
 	}
 }
 interface Env extends Cloudflare.Env {}

apps/cloudflare-one-casb/wrangler.jsonc

@@ -18,8 +18,11 @@
 	},
 	"vars": {
 		"MCP_SERVER_NAME": "PLACEHOLDER",
-		"MCP_SERVER_VERSION": "PLACEHOLDER"
+		"MCP_SERVER_VERSION": "PLACEHOLDER",
+		"CLOUDFLARE_CLIENT_ID": "<PLACEHOLDER>",
+		"CLOUDFLARE_CLIENT_SECRET": "<PLACEHOLDER>"
 	},
+
 	"dev": {
 		"port": 8976
 	},

packages/mcp-common/src/schemas/cf1-integrations.ts

@@ -73,7 +73,6 @@ export type zReturnedAssetsResult = z.infer<typeof AssetsResponse>
 export const AssetCategoriesResponse = z.array(AssetCategory)
 export type zReturnedAssetCategoriesResult = z.infer<typeof AssetCategoriesResponse>
 
-
 export const assetCategoryTypeParam = z
 	.enum([
 		'Account',
@@ -127,7 +126,7 @@ export const assetCategoryTypeParam = z
 	.optional()
 	.describe('Type of cloud resource or service category')
 
-	export const assetCategoryVendorParam = z
+export const assetCategoryVendorParam = z
 	.enum([
 		'AWS',
 		'Bitbucket',
@@ -147,4 +146,4 @@ export const assetCategoryTypeParam = z
 		'Workday',
 		'Zoom',
 	])
-	.describe('Vendor of the cloud service or resource')
+	.describe('Vendor of the cloud service or resource')