diff --git a/cmd/api-server/main.go b/cmd/api-server/main.go
index a21e36a8ed9..f110d0f1d9d 100644
--- a/cmd/api-server/main.go
+++ b/cmd/api-server/main.go
@@ -328,11 +328,11 @@ func main() {
 	getTestWorkflowServiceNotificationsStream := func(ctx context.Context, executionID, serviceName string, serviceIndex int) (<-chan testkube.TestWorkflowExecutionNotification, error) {
 		execution, err := testWorkflowResultsRepository.Get(ctx, executionID)
 		if err != nil {
-			return nil, errors.Join(err, agent.ErrGetTestWorkflowExecution)
+			return nil, err
 		}
 
 		if execution.Result != nil && execution.Result.IsFinished() {
-			return nil, agent.ErrFinishedTestWorkflowExecution
+			return nil, errors.New("test workflow execution is finished")
 		}
 
 		notifications := executionWorker.Notifications(ctx, fmt.Sprintf("%s-%s-%d", execution.Id, serviceName, serviceIndex), executionworkertypes.NotificationsOptions{
diff --git a/pkg/agent/testworkflows.go b/pkg/agent/testworkflows.go
index 65fec85ca17..8924284b506 100644
--- a/pkg/agent/testworkflows.go
+++ b/pkg/agent/testworkflows.go
@@ -16,15 +16,11 @@ import (
 	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
 	"github.com/kubeshop/testkube/pkg/cloud"
 	"github.com/kubeshop/testkube/pkg/testworkflows/executionworker/controller"
+	"github.com/kubeshop/testkube/pkg/testworkflows/executionworker/registry"
 )
 
 const testWorkflowNotificationsRetryCount = 10
 
-var (
-	ErrGetTestWorkflowExecution      = errors.New("can't get test workflow execution")
-	ErrFinishedTestWorkflowExecution = errors.New("test workflow execution is finished")
-)
-
 func getTestWorkflowNotificationType(n testkube.TestWorkflowExecutionNotification) cloud.TestWorkflowNotificationType {
 	if n.Result != nil {
 		return cloud.TestWorkflowNotificationType_WORKFLOW_STREAM_RESULT
@@ -248,11 +244,7 @@ func (ag *Agent) executeWorkflowServiceNotificationsRequest(ctx context.Context,
 
 	for {
 		notificationsCh, err = ag.testWorkflowServiceNotificationsFunc(ctx, req.ExecutionId, req.ServiceName, int(req.ServiceIndex))
-		if errors.Is(err, ErrGetTestWorkflowExecution) || errors.Is(err, ErrFinishedTestWorkflowExecution) {
-			break
-		}
-
-		if err != nil {
+		if errors.Is(err, registry.ErrResourceNotFound) {
 			// We have a race condition here
 			// Cloud sometimes slow to start service
 			// while WorkflowNotifications request from websockets comes in faster