diff --git a/packages/cli/src/evaluation/test-definitions.service.ee.ts b/packages/cli/src/evaluation/test-definitions.service.ee.ts
index 058c02d8f6..e9b1b0da09 100644
--- a/packages/cli/src/evaluation/test-definitions.service.ee.ts
+++ b/packages/cli/src/evaluation/test-definitions.service.ee.ts
@@ -1,10 +1,12 @@
 import { Service } from 'typedi';
 
 import type { TestDefinition } from '@/databases/entities/test-definition.ee';
+import { AnnotationTagRepository } from '@/databases/repositories/annotation-tag.repository.ee';
 import { TestDefinitionRepository } from '@/databases/repositories/test-definition.repository.ee';
 import { NotFoundError } from '@/errors/response-errors/not-found.error';
 import { validateEntity } from '@/generic-helpers';
 import type { ListQuery } from '@/requests';
+import { BadRequestError } from '@/errors/response-errors/bad-request.error';
 
 type TestDefinitionLike = Omit<
 	Partial<TestDefinition>,
@@ -17,7 +19,10 @@ type TestDefinitionLike = Omit<
 
 @Service()
 export class TestDefinitionsService {
-	constructor(private testRepository: TestDefinitionRepository) {}
+	constructor(
+		private testRepository: TestDefinitionRepository,
+		private annotationTagRepository: AnnotationTagRepository,
+	) {}
 
 	private toEntityLike(attrs: {
 		name?: string;
@@ -89,6 +94,20 @@ export class TestDefinitionsService {
 			delete attrs.workflowId;
 		}
 
+		// Check if the annotation tag exists
+		if (attrs.annotationTagId) {
+			const annotationTagExists = await this.annotationTagRepository.exists({
+				where: {
+					id: attrs.annotationTagId,
+				},
+			});
+
+			if (!annotationTagExists) {
+				throw new BadRequestError('Annotation tag not found');
+			}
+		}
+
+		// Update the test definition
 		const queryResult = await this.testRepository.update(id, this.toEntityLike(attrs));
 
 		if (queryResult.affected === 0) {
diff --git a/packages/cli/test/integration/evaluation/test-definitions.api.test.ts b/packages/cli/test/integration/evaluation/test-definitions.api.test.ts
index 5d75b21ffd..7fffad1a38 100644
--- a/packages/cli/test/integration/evaluation/test-definitions.api.test.ts
+++ b/packages/cli/test/integration/evaluation/test-definitions.api.test.ts
@@ -11,6 +11,8 @@ import { createWorkflow } from './../shared/db/workflows';
 import * as testDb from './../shared/test-db';
 import type { SuperAgentTest } from './../shared/types';
 import * as utils from './../shared/utils/';
+import { AnnotationTagEntity } from '@/databases/entities/annotation-tag-entity.ee';
+import { createAnnotationTags } from '@test-integration/db/executions';
 
 let authOwnerAgent: SuperAgentTest;
 let workflowUnderTest: WorkflowEntity;
@@ -358,4 +360,34 @@ describe('DELETE /evaluation/test-definitions/:id', () => {
 		expect(resp.statusCode).toBe(404);
 		expect(resp.body.message).toBe('Test definition not found');
 	});
+
+	test('should update annotationTagId', async () => {
+		const newTest = Container.get(TestDefinitionRepository).create({
+			name: 'test',
+			workflow: { id: workflowUnderTest.id },
+		});
+		await Container.get(TestDefinitionRepository).save(newTest);
+
+		const resp = await authOwnerAgent.patch(`/evaluation/test-definitions/${newTest.id}`).send({
+			annotationTagId: annotationTag.id,
+		});
+
+		expect(resp.statusCode).toBe(200);
+		expect(resp.body.data.annotationTag.id).toBe(annotationTag.id);
+	});
+
+	test('should return error if annotationTagId is invalid', async () => {
+		const newTest = Container.get(TestDefinitionRepository).create({
+			name: 'test',
+			workflow: { id: workflowUnderTest.id },
+		});
+		await Container.get(TestDefinitionRepository).save(newTest);
+
+		const resp = await authOwnerAgent.patch(`/evaluation/test-definitions/${newTest.id}`).send({
+			annotationTagId: 123,
+		});
+
+		expect(resp.statusCode).toBe(400);
+		expect(resp.body.message).toBe('Annotation tag not found');
+	});
 });