Integration tests stabilization (#559)

Author: aozherelyeva

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/AIAgentIntegrationTest.kt

@@ -981,14 +981,13 @@ class AIAgentIntegrationTest {
     @MethodSource("openAIModels", "anthropicModels", "googleModels")
     fun integration_AgentWithToolsWithoutParamsTest(model: LLModel) = runTest(timeout = 120.seconds) {
         assumeTrue(model.capabilities.contains(LLMCapability.Tools), "Model $model does not support tools")
-        assumeTrue(
-            model.id != GoogleModels.Gemini2_0Flash.id,
-            "gemini-2.0-flash returns flaky results and fails to call tools on a permanent basis"
-        )
-        assumeTrue(
-            model.id != GoogleModels.Gemini2_0Flash001.id,
-            "gemini-2.0-flash-001 returns flaky results and fails to call tools on a permanent basis"
+        val flakyModels = listOf(
+            GoogleModels.Gemini2_0Flash.id,
+            GoogleModels.Gemini2_0Flash001.id,
+            GoogleModels.Gemini2_0FlashLite.id,
+            GoogleModels.Gemini2_0FlashLite001.id
         )
+        assumeTrue(!flakyModels.contains(model.id), "Model $model is flaky and fails to call tools")
 
         val registry = ToolRegistry {
             tool(CalculatorToolNoArgs)

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/SingleLLMPromptExecutorIntegrationTest.kt

@@ -716,32 +716,34 @@ class SingleLLMPromptExecutorIntegrationTest {
                     }
                 }
 
-            try {
-                val response = executor.execute(prompt, model)
-                when (scenario) {
-                    MarkdownTestScenario.MALFORMED_SYNTAX,
-                    MarkdownTestScenario.MATH_NOTATION,
-                    MarkdownTestScenario.BROKEN_LINKS,
-                    MarkdownTestScenario.IRREGULAR_TABLES -> {
-                        checkResponseBasic(response)
-                    }
+            withRetry {
+                try {
+                    val response = executor.execute(prompt, model)
+                    when (scenario) {
+                        MarkdownTestScenario.MALFORMED_SYNTAX,
+                        MarkdownTestScenario.MATH_NOTATION,
+                        MarkdownTestScenario.BROKEN_LINKS,
+                        MarkdownTestScenario.IRREGULAR_TABLES -> {
+                            checkResponseBasic(response)
+                        }
 
-                    else -> {
-                        checkExecutorMediaResponse(response)
+                        else -> {
+                            checkExecutorMediaResponse(response)
+                        }
                     }
-                }
-            } catch (e: Exception) {
-                when (scenario) {
-                    MarkdownTestScenario.EMPTY_MARKDOWN -> {
-                        when (model.provider) {
-                            LLMProvider.Google -> {
-                                println("Expected exception for ${scenario.name.lowercase()} image: ${e.message}")
+                } catch (e: Exception) {
+                    when (scenario) {
+                        MarkdownTestScenario.EMPTY_MARKDOWN -> {
+                            when (model.provider) {
+                                LLMProvider.Google -> {
+                                    println("Expected exception for ${scenario.name.lowercase()} image: ${e.message}")
+                                }
                             }
                         }
-                    }
 
-                    else -> {
-                        throw e
+                        else -> {
+                            throw e
+                        }
                     }
                 }
             }