(feat): Add granular user agent for Ai (#13835)

Author: yuhengshs

packages/api-graphql/__tests__/GraphQLAPI.test.ts

@@ -6,6 +6,8 @@ import * as typedQueries from './fixtures/with-types/queries';
 import * as typedSubscriptions from './fixtures/with-types/subscriptions';
 import { expectGet } from './utils/expects';
 import { InternalGraphQLAPIClass } from '../src/internals/InternalGraphQLAPI';
+import { GraphQLAuthMode } from '@aws-amplify/core/internals/utils';
+import { INTERNAL_USER_AGENT_OVERRIDE } from '@aws-amplify/data-schema/runtime';
 
 import {
 	__amplify,
@@ -1614,4 +1616,58 @@ describe('API test', () => {
 		const subscribeOptions = spyon_appsync_realtime.mock.calls[0][0];
 		expect(subscribeOptions).toBe(resolvedUrl);
 	});
+	test('graphql method handles INTERNAL_USER_AGENT_OVERRIDE correctly', async () => {
+		Amplify.configure({
+			API: {
+				GraphQL: {
+					defaultAuthMode: 'apiKey',
+					apiKey: 'FAKE-KEY',
+					endpoint: 'https://localhost/graphql',
+					region: 'local-host-h4x',
+				},
+			},
+		});
+
+		const mockPost = jest.fn().mockResolvedValue({
+			body: {
+				json: () => ({ data: { test: 'result' } }),
+			},
+		});
+		(raw.GraphQLAPI as any)._api.post = mockPost;
+
+		const graphqlOptions = {
+			query: 'query TestQuery { test }',
+			variables: { id: 'some-id' },
+			authMode: 'apiKey' as GraphQLAuthMode,
+			[INTERNAL_USER_AGENT_OVERRIDE]: {
+				category: 'CustomCategory',
+				action: 'CustomAction',
+			},
+		};
+
+		await client.graphql(graphqlOptions);
+
+		// Check if the INTERNAL_USER_AGENT_OVERRIDE was properly handled
+		expect(mockPost).toHaveBeenCalledWith(
+			expect.anything(),
+			expect.objectContaining({
+				options: expect.objectContaining({
+					headers: expect.objectContaining({
+						'x-amz-user-agent': expect.stringContaining(
+							'CustomCategory/CustomAction',
+						),
+					}),
+				}),
+			}),
+		);
+		// Ensure the INTERNAL_USER_AGENT_OVERRIDE was not passed along in the options
+		expect(mockPost).not.toHaveBeenCalledWith(
+			expect.anything(),
+			expect.objectContaining({
+				options: expect.objectContaining({
+					[INTERNAL_USER_AGENT_OVERRIDE]: expect.anything(),
+				}),
+			}),
+		);
+	});
 });

packages/api-graphql/package.json

@@ -84,9 +84,12 @@
 		"server"
 	],
 	"dependencies": {
+		"@aws-amplify/api-rest": "4.0.49",
+		"@aws-amplify/core": "6.4.2",
+		"@aws-amplify/data-schema": "^1.7.0",
 		"@aws-amplify/api-rest": "4.0.50",
 		"@aws-amplify/core": "6.4.3",
-		"@aws-amplify/data-schema": "^1.5.0",
+		"@aws-amplify/data-schema": "^1.7.0",
 		"@aws-sdk/types": "3.387.0",
 		"graphql": "15.8.0",
 		"rxjs": "^7.8.1",

packages/api-graphql/src/GraphQLAPI.ts

@@ -1,13 +1,28 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 import { AmplifyClassV6 } from '@aws-amplify/core';
-import { ApiAction, Category } from '@aws-amplify/core/internals/utils';
-import { CustomHeaders } from '@aws-amplify/data-schema/runtime';
+import {
+	ApiAction,
+	Category,
+	CustomUserAgentDetails,
+} from '@aws-amplify/core/internals/utils';
+import {
+	CustomHeaders,
+	INTERNAL_USER_AGENT_OVERRIDE,
+} from '@aws-amplify/data-schema/runtime';
 import { Observable } from 'rxjs';
 
 import { GraphQLOptions, GraphQLResult } from './types';
 import { InternalGraphQLAPIClass } from './internals/InternalGraphQLAPI';
 
+function isGraphQLOptionsWithOverride(
+	options: GraphQLOptions,
+): options is GraphQLOptions & {
+	[INTERNAL_USER_AGENT_OVERRIDE]: CustomUserAgentDetails;
+} {
+	return INTERNAL_USER_AGENT_OVERRIDE in options;
+}
+
 export const graphqlOperation = (
 	query: any,
 	variables = {},
@@ -38,9 +53,25 @@ export class GraphQLAPIClass extends InternalGraphQLAPIClass {
 		options: GraphQLOptions,
 		additionalHeaders?: CustomHeaders,
 	): Observable<GraphQLResult<T>> | Promise<GraphQLResult<T>> {
-		return super.graphql(amplify, options, additionalHeaders, {
+		const userAgentDetails: CustomUserAgentDetails = {
 			category: Category.API,
 			action: ApiAction.GraphQl,
+		};
+
+		if (isGraphQLOptionsWithOverride(options)) {
+			const {
+				[INTERNAL_USER_AGENT_OVERRIDE]: internalUserAgentOverride,
+				...cleanOptions
+			} = options;
+
+			return super.graphql(amplify, cleanOptions, additionalHeaders, {
+				...userAgentDetails,
+				...internalUserAgentOverride,
+			});
+		}
+
+		return super.graphql(amplify, options, additionalHeaders, {
+			...userAgentDetails,
 		});
 	}
 

packages/aws-amplify/package.json

@@ -293,7 +293,7 @@
 			"name": "[Analytics] record (Pinpoint)",
 			"path": "./dist/esm/analytics/index.mjs",
 			"import": "{ record }",
-			"limit": "17.25 kB"
+			"limit": "17.35 kB"
 		},
 		{
 			"name": "[Analytics] record (Kinesis)",
@@ -317,7 +317,7 @@
 			"name": "[Analytics] identifyUser (Pinpoint)",
 			"path": "./dist/esm/analytics/index.mjs",
 			"import": "{ identifyUser }",
-			"limit": "15.75 kB"
+			"limit": "15.85 kB"
 		},
 		{
 			"name": "[Analytics] enable",
@@ -335,7 +335,7 @@
 			"name": "[API] generateClient (AppSync)",
 			"path": "./dist/esm/api/index.mjs",
 			"import": "{ generateClient }",
-			"limit": "43.1 kB"
+			"limit": "43.40 kB"
 		},
 		{
 			"name": "[API] REST API handlers",
@@ -461,43 +461,43 @@
 			"name": "[Storage] copy (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ copy }",
-			"limit": "14.86 kB"
+			"limit": "14.96 kB"
 		},
 		{
 			"name": "[Storage] downloadData (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ downloadData }",
-			"limit": "15.45 kB"
+			"limit": "15.55 kB"
 		},
 		{
 			"name": "[Storage] getProperties (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ getProperties }",
-			"limit": "14.71 kB"
+			"limit": "14.81 kB"
 		},
 		{
 			"name": "[Storage] getUrl (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ getUrl }",
-			"limit": "15.95 kB"
+			"limit": "16.05 kB"
 		},
 		{
 			"name": "[Storage] list (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ list }",
-			"limit": "15.31 kB"
+			"limit": "15.41 kB"
 		},
 		{
 			"name": "[Storage] remove (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ remove }",
-			"limit": "14.57 kB"
+			"limit": "14.67 kB"
 		},
 		{
 			"name": "[Storage] uploadData (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ uploadData }",
-			"limit": "19.95 kB"
+			"limit": "20.05 kB"
 		}
 	]
 }

packages/core/src/Platform/types.ts

@@ -26,6 +26,7 @@ export enum Framework {
 }
 
 export enum Category {
+	AI = 'ai',
 	API = 'api',
 	Auth = 'auth',
 	Analytics = 'analytics',
@@ -39,6 +40,17 @@ export enum Category {
 	Storage = 'storage',
 }
 
+export enum AiAction {
+	CreateConversation = '1',
+	GetConversation = '2',
+	ListConversations = '3',
+	DeleteConversation = '4',
+	SendMessage = '5',
+	ListMessages = '6',
+	OnMessage = '7',
+	Generation = '8',
+}
+
 export enum AnalyticsAction {
 	Record = '1',
 	IdentifyUser = '2',
@@ -123,6 +135,7 @@ export enum StorageAction {
 }
 
 interface ActionMap {
+	[Category.AI]: AiAction;
 	[Category.Auth]: AuthAction;
 	[Category.API]: ApiAction;
 	[Category.Analytics]: AnalyticsAction;
@@ -148,6 +161,7 @@ interface CustomUserAgentDetailsBase {
 
 export type CustomUserAgentDetails =
 	| (CustomUserAgentDetailsBase & { category?: never; action?: never })
+	| UserAgentDetailsWithCategory<Category.AI>
 	| UserAgentDetailsWithCategory<Category.API>
 	| UserAgentDetailsWithCategory<Category.Auth>
 	| UserAgentDetailsWithCategory<Category.Analytics>
@@ -180,6 +194,12 @@ export interface StorageUserAgentInput {
 	additionalDetails: AdditionalDetails;
 }
 
+export interface AiUserAgentInput {
+	category: Category.AI;
+	apis: AiAction[];
+	additionalDetails: AdditionalDetails;
+}
+
 export interface AuthUserAgentInput {
 	category: Category.Auth;
 	apis: AuthAction[];
@@ -202,4 +222,5 @@ export type SetCustomUserAgentInput =
 	| StorageUserAgentInput
 	| AuthUserAgentInput
 	| InAppMessagingUserAgentInput
-	| GeoUserAgentInput;
+	| GeoUserAgentInput
+	| AiUserAgentInput;

packages/core/src/libraryUtils.ts

@@ -78,6 +78,7 @@ export {
 	getAmplifyUserAgent,
 } from './Platform';
 export {
+	AiAction,
 	ApiAction,
 	AuthAction,
 	AnalyticsAction,