Add @Retry annotation (#230)

Author: aozherelyeva

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

@@ -1,9 +1,9 @@
 package ai.koog.integration.tests
 
-import ai.koog.integration.tests.utils.TestUtils.executeWithRetry
 import ai.koog.integration.tests.utils.TestUtils.readTestAnthropicKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestGoogleAIKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestOpenAIKeyFromEnv
+import ai.koog.integration.tests.utils.annotations.Retry
 import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.model.PromptExecutorExt.execute
 import ai.koog.prompt.executor.clients.anthropic.AnthropicLLMClient
@@ -23,6 +23,7 @@ class MultipleSystemMessagesPromptIntegrationTest {
     private val anthropicApiKey = readTestAnthropicKeyFromEnv()
     private val googleApiKey = readTestGoogleAIKeyFromEnv()
 
+    @Retry(3)
     @Test
     fun integration_testMultipleSystemMessages() = runBlocking {
         val openAIClient = OpenAILLMClient(openAIApiKey)
@@ -46,9 +47,9 @@ class MultipleSystemMessagesPromptIntegrationTest {
         val modelAnthropic = AnthropicModels.Haiku_3_5
         val modelGemini = GoogleModels.Gemini2_0Flash
 
-        val responseOpenAI = executeWithRetry { executor.execute(prompt, modelOpenAI) }
-        val responseAnthropic = executeWithRetry { executor.execute(prompt, modelAnthropic) }
-        val responseGemini = executeWithRetry { executor.execute(prompt, modelGemini) }
+        val responseOpenAI = executor.execute(prompt, modelOpenAI)
+        val responseAnthropic = executor.execute(prompt, modelAnthropic)
+        val responseGemini = executor.execute(prompt, modelGemini)
 
         assertTrue(responseOpenAI.content.isNotEmpty(), "OpenAI response should not be empty")
         assertTrue(responseAnthropic.content.isNotEmpty(), "Anthropic response should not be empty")

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

@@ -13,7 +13,7 @@ import ai.koog.agents.core.dsl.builder.strategy
 import ai.koog.agents.core.dsl.extension.*
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.agents.features.eventHandler.feature.EventHandler
-import ai.koog.integration.tests.utils.TestUtils.runWithRetry
+import ai.koog.integration.tests.utils.annotations.Retry
 import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.OllamaModels
@@ -174,13 +174,14 @@ class OllamaAgentIntegrationTest {
         }
     }
 
+    @Retry(3)
     @Test
     fun ollama_testAgentClearContext() = runTest(timeout = 600.seconds) {
         val strategy = createTestStrategy()
         val toolRegistry = createToolRegistry()
         val agent = createAgent(executor, strategy, toolRegistry)
 
-        val result = runWithRetry { agent.runAndGetResult("What is the capital of France?") }
+        val result = agent.runAndGetResult("What is the capital of France?")
 
         assertNotNull(result, "Result should not be empty")
         assertTrue(result.isNotEmpty(), "Result should not be empty")

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

@@ -5,7 +5,7 @@ import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.agents.ext.tool.SayToUser
 import ai.koog.agents.features.eventHandler.feature.EventHandler
 import ai.koog.agents.features.eventHandler.feature.EventHandlerConfig
-import ai.koog.integration.tests.utils.TestUtils.runWithRetry
+import ai.koog.integration.tests.utils.annotations.Retry
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.extension.ExtendWith
 import kotlin.test.AfterTest
@@ -87,6 +87,7 @@ class OllamaSimpleAgentIntegrationTest {
         actualToolCalls.clear()
     }
 
+    @Retry(3)
     @Test
     fun ollama_simpleTest() = runTest(timeout = 600.seconds) {
         val toolRegistry = ToolRegistry.Companion {
@@ -119,9 +120,8 @@ class OllamaSimpleAgentIntegrationTest {
             installFeatures = { install(EventHandler.Feature, eventHandlerConfig) }
         )
 
-        runWithRetry {
-            agent.run("Give me top 10 books of the all time.")
-        }
+
+        agent.run("Give me top 10 books of the all time.")
 
         assertTrue(actualToolCalls.isNotEmpty(), "No tools were called for model")
         assertTrue(

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

@@ -9,7 +9,7 @@ import ai.koog.integration.tests.utils.TestUtils.CalculatorTool
 import ai.koog.integration.tests.utils.TestUtils.readTestAnthropicKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestGoogleAIKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestOpenAIKeyFromEnv
-import ai.koog.integration.tests.utils.TestUtils.runWithRetry
+import ai.koog.integration.tests.utils.annotations.Retry
 import ai.koog.prompt.executor.clients.google.GoogleModels
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.executor.llms.all.simpleAnthropicExecutor
@@ -116,6 +116,7 @@ class SimpleAgentIntegrationTest {
         results.clear()
     }
 
+    @Retry(3)
     @ParameterizedTest
     @MethodSource("openAIModels", "anthropicModels", "googleModels")
     fun integration_AIAgentShouldNotCallToolsByDefault(model: LLModel) = runBlocking {
@@ -134,14 +135,13 @@ class SimpleAgentIntegrationTest {
             installFeatures = { install(EventHandler.Feature, eventHandlerConfig) }
         )
 
-        runWithRetry {
-            agent.run("Repeat what I say: hello, I'm good.")
-        }
+        agent.run("Repeat what I say: hello, I'm good.")
 
         // by default, AIAgent has no tools underneath
         assertTrue(actualToolCalls.isEmpty(), "No tools should be called for model $model")
     }
 
+    @Retry(3)
     @ParameterizedTest
     @MethodSource("openAIModels", "anthropicModels", "googleModels")
     fun integration_AIAgentShouldCallCustomTool(model: LLModel) = runBlocking {
@@ -172,9 +172,8 @@ class SimpleAgentIntegrationTest {
             installFeatures = { install(EventHandler.Feature, eventHandlerConfig) }
         )
 
-        runWithRetry {
-            agent.run("How much is 3 times 5?")
-        }
+
+        agent.run("How much is 3 times 5?")
 
         assertTrue(actualToolCalls.isNotEmpty(), "No tools were called for model $model")
         assertTrue(