From 5c8d39d7dff89284da9798d7447847d0362b7d86 Mon Sep 17 00:00:00 2001 From: Erik Dix Date: Sat, 6 Dec 2025 14:53:20 +0100 Subject: [PATCH] test --- .gitignore | 2 + pom.xml | 31 ++- .../softwarefabric/agents/AgenTester2.kt | 143 +++++++++++++ .../softwarefabric/agents/AgentTester.kt | 188 ++++++++++++++++++ .../softwarefabric/agents/DesignerAgent.kt | 8 + .../softwarefabric/agents/SnapshotStrategy.kt | 53 +++++ .../softwarefabric/client/TestAgent.kt | 59 ++++++ .../controllers/A2AController2.kt | 166 ++++++++++++++++ .../storage/FileMessageStorage.kt | 172 ++++++++++++++++ 9 files changed, 819 insertions(+), 3 deletions(-) create mode 100644 src/main/kotlin/nl/trivion/softwarefabric/agents/AgenTester2.kt create mode 100644 src/main/kotlin/nl/trivion/softwarefabric/agents/AgentTester.kt create mode 100644 src/main/kotlin/nl/trivion/softwarefabric/agents/SnapshotStrategy.kt create mode 100644 src/main/kotlin/nl/trivion/softwarefabric/client/TestAgent.kt create mode 100644 src/main/kotlin/nl/trivion/softwarefabric/controllers/A2AController2.kt create mode 100644 src/main/kotlin/nl/trivion/softwarefabric/storage/FileMessageStorage.kt diff --git a/.gitignore b/.gitignore index 2485648..2c2ec89 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,5 @@ build/ ### Mac OS ### .DS_Store /.idea/ +/agent-persistence-test/ +/a2a-messages/ diff --git a/pom.xml b/pom.xml index 91b9924..ece7df7 100644 --- a/pom.xml +++ b/pom.xml @@ -16,6 +16,7 @@ 21 2.2.10 + 0.5.3 @@ -50,6 +51,10 @@ org.springframework.boot spring-boot-starter + + org.springframework.boot + spring-boot-starter-web + org.jetbrains.kotlin kotlin-reflect @@ -63,17 +68,37 @@ ai.koog koog-agents-jvm - 0.5.0 + ${koog.version} ai.koog prompt-cache-redis-jvm - 0.5.0 + ${koog.version} ai.koog prompt-executor-model-jvm - 0.5.0 + ${koog.version} + + + ai.koog + agents-features-a2a-client-jvm + ${koog.version} + + + ai.koog + agents-features-a2a-server-jvm + ${koog.version} + + + ai.koog + a2a-transport-server-jsonrpc-http-jvm + ${koog.version} + + + ai.koog + a2a-transport-client-jsonrpc-http-jvm + ${koog.version} diff --git a/src/main/kotlin/nl/trivion/softwarefabric/agents/AgenTester2.kt b/src/main/kotlin/nl/trivion/softwarefabric/agents/AgenTester2.kt new file mode 100644 index 0000000..7ea82fa --- /dev/null +++ b/src/main/kotlin/nl/trivion/softwarefabric/agents/AgenTester2.kt @@ -0,0 +1,143 @@ +package nl.trivion.softwarefabric.agents + +import ai.koog.agents.core.agent.AIAgent +import ai.koog.agents.core.agent.config.AIAgentConfig +import ai.koog.agents.core.agent.singleRunStrategy +import ai.koog.agents.core.dsl.builder.strategy +import ai.koog.agents.core.tools.ToolRegistry +import ai.koog.agents.ext.tool.AskUser +import ai.koog.agents.ext.tool.SayToUser +import ai.koog.agents.snapshot.feature.Persistence +import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import ai.koog.prompt.dsl.prompt +import ai.koog.prompt.executor.llms.all.simpleOllamaAIExecutor +import ai.koog.prompt.llm.LLMCapability +import ai.koog.prompt.llm.LLMProvider +import ai.koog.prompt.llm.LLModel +import ai.koog.prompt.llm.OllamaModels +import kotlinx.coroutines.runBlocking +import java.nio.file.Files +import kotlin.io.path.exists +import kotlin.uuid.ExperimentalUuidApi + +/** + * This example demonstrates how to use the file-based checkpoint provider with a persistent agent. + * + * The JVMFileAgentCheckpointStorageProvider stores agent checkpoints in a file system, + * allowing the agent's state to persist across multiple runs. + * + * This example shows: + * 1. How to create a file-based checkpoint provider + * 2. How to configure an agent with the file-based checkpoint provider + * 3. How to run an agent that automatically creates checkpoints + * 4. How to restore an agent from a checkpoint + */ +@OptIn(ExperimentalUuidApi::class) +fun main() = runBlocking { + // Create a temporary directory for storing checkpoints + val checkpointDir = Files.createTempDirectory("agent-checkpoints") + println("Checkpoint directory: $checkpointDir") + + // Create the file-based checkpoint provider + val provider = JVMFilePersistenceStorageProvider(checkpointDir) + + // Create a unique agent ID to identify this agent's checkpoints + val agentId = "persistent-agent-example3" + + // Create tool registry with basic tools + val toolRegistry = ToolRegistry { + tool(AskUser) + tool(SayToUser) + } + + val customModel: LLModel = LLModel( + provider = LLMProvider.Ollama, + id = "gpt-oss:20b", + capabilities = listOf( + LLMCapability.Temperature + ), + contextLength = 2048 + ) + + // Create agent config with a system prompt + val agentConfig = AIAgentConfig( + prompt = prompt("persistent-agent") { + system("You are a helpful assistant that remembers conversations across sessions.") + }, + model = customModel, + maxAgentIterations = 20 + ) + + println("Creating and running agent with continuous persistence...") + + simpleOllamaAIExecutor().use { executor -> + // Create the agent with the file-based checkpoint provider and continuous persistence + val agent = AIAgent( + promptExecutor = executor, +// strategy = SnapshotStrategy.strategy, + strategy = singleRunStrategy(), + agentConfig = agentConfig, + toolRegistry = toolRegistry, + id = agentId + ) { + install(Persistence) { + storage = provider // Use the file-based checkpoint provider + enableAutomaticPersistence = true // Enable automatic checkpoint creation + } + } + + // Run the agent with an initial input + val result = agent.run("Hallo ik heet Erik") + println("Agent result: $result") + + // Retrieve all checkpoints created during the agent's execution + val checkpoints = provider.getCheckpoints(agentId) + println("\nRetrieved ${checkpoints.size} checkpoints for agent $agentId") + + // Print checkpoint details + checkpoints.forEachIndexed { index, checkpoint -> + println("Checkpoint ${index + 1}:") + println(" ID: ${checkpoint.checkpointId}") + println(" Created at: ${checkpoint.createdAt}") + println(" Node ID: ${checkpoint.nodeId}") + println(" Message history size: ${checkpoint.messageHistory.size}") + } + + // Verify that the checkpoint files exist in the file system + val checkpointsDir = checkpointDir.resolve("checkpoints").resolve(agentId) + if (checkpointsDir.exists()) { + println("\nCheckpoint files in directory: ${checkpointsDir.toFile().list()?.joinToString()}") + } + + println("\nNow creating a new agent instance with the same ID to demonstrate restoration...") + + // Create a new agent instance with the same ID + // It will automatically restore from the latest checkpoint + val restoredAgent = AIAgent( + promptExecutor = executor, +// strategy = SnapshotStrategy.strategy, + strategy = singleRunStrategy(), + agentConfig = agentConfig, + toolRegistry = toolRegistry, + id = agentId + ) { + install(Persistence) { + storage = provider // Use the file-based checkpoint provider + enableAutomaticPersistence = true // Enable automatic checkpoint creation + } + } + + // Run the restored agent with a new input + // The agent will continue the conversation from where it left off + val restoredResult = restoredAgent.run("Hoe heet ik?") + println("Restored agent result: $restoredResult") + + // Get the latest checkpoint after the second run + val latestCheckpoint = provider.getLatestCheckpoint(agentId) + println("\nLatest checkpoint after restoration:") + println(" ID: ${latestCheckpoint?.checkpointId}") + println(" Created at: ${latestCheckpoint?.createdAt}") + println(" Node ID: ${latestCheckpoint?.nodeId}") + println(" Message history size: ${latestCheckpoint?.messageHistory?.size}") + } +} \ No newline at end of file diff --git a/src/main/kotlin/nl/trivion/softwarefabric/agents/AgentTester.kt b/src/main/kotlin/nl/trivion/softwarefabric/agents/AgentTester.kt new file mode 100644 index 0000000..1657ac1 --- /dev/null +++ b/src/main/kotlin/nl/trivion/softwarefabric/agents/AgentTester.kt @@ -0,0 +1,188 @@ +package nl.trivion.softwarefabric.agents + +import ai.koog.agents.core.agent.AIAgent +import ai.koog.agents.core.agent.config.AIAgentConfig +import ai.koog.agents.core.dsl.builder.forwardTo +import ai.koog.agents.core.dsl.builder.strategy +import ai.koog.agents.core.dsl.extension.nodeLLMRequest +import ai.koog.agents.features.eventHandler.feature.EventHandler +import ai.koog.agents.snapshot.feature.Persistence +import ai.koog.agents.snapshot.providers.InMemoryPersistenceStorageProvider +import ai.koog.agents.snapshot.providers.PersistenceStorageProvider +import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.executor.llms.all.simpleOllamaAIExecutor +import ai.koog.prompt.llm.LLMCapability +import ai.koog.prompt.llm.LLMProvider +import ai.koog.prompt.llm.LLModel +import kotlinx.coroutines.runBlocking +import org.slf4j.LoggerFactory +import java.io.File + +class AgentTester { + + companion object { + private val log = LoggerFactory.getLogger(AgentTester::class.java) + + @JvmStatic + fun main(args: Array) { + println("=== Simple Agent Persistence Test ===") + println() + + // Parse command line arguments + val issueNumber = args.getOrNull(0)?.toIntOrNull() ?: 1 + val message = args.getOrNull(1) ?: "Ik heet Erik." + + println("Testing with:") + println("- Issue number: $issueNumber") + println("- Message: $message") + println() + + val tester = AgentTester() + + runBlocking { + tester.testAgent(issueNumber, message) + } + } + } + + private val promptExecutor = simpleOllamaAIExecutor("http://localhost:11434") + + private val customModel: LLModel = LLModel( + provider = LLMProvider.Ollama, + id = "gpt-oss:20b", + capabilities = listOf( + LLMCapability.Temperature + ), + contextLength = 2048 + ) + + private val systemPrompt = """ + You are a helpful assistant. Keep your responses concise and friendly. + """.trimIndent() + + private val agentConfig = AIAgentConfig( + prompt = Prompt.build("simple-chat") { + system(systemPrompt) + }, + model = customModel, + maxAgentIterations = 5 + ) + + private fun createAgent(agentId: String, provider: PersistenceStorageProvider<*>): AIAgent { + + return AIAgent( + promptExecutor = promptExecutor, + agentConfig = agentConfig, + id = agentId, + installFeatures = { + install(EventHandler) { + onAgentCompleted { result -> + log.info("✓ Agent completed") + log.info(" Result: ${result.result}") + + } + } + + install(Persistence) { + storage = provider + enableAutomaticPersistence = true + + log.info("✓ Persistence installed with automatic persistence enabled") + } + } + ) + } + + suspend fun testAgent(issueNumber: Int, message: String) { + val agentId = "developer-agent-issue-$issueNumber" + + println("Creating agent with ID: $agentId") + println("-".repeat(50)) + + try { + val persistenceDir = File("./agent-persistence-test") + + if (!persistenceDir.exists()) { + persistenceDir.mkdirs() + log.info("Created persistence directory: ${persistenceDir.absolutePath}") + } + +// val provider = JVMFilePersistenceStorageProvider(persistenceDir.toPath()) + val provider = InMemoryPersistenceStorageProvider() + + log.info("Creating agent with ID: $agentId") + log.info("Persistence directory: ${persistenceDir.absolutePath}") + + var agent = createAgent(agentId, provider) + + println() + println("Running agent with message: '$message'") + println("-".repeat(50)) + + var result = agent.run(message) + + println() + println("=== RESULT ===") + println(result) + println() + + agent = createAgent(agentId, provider) + result = agent.run("Hoe heet ik?") + agent.close() + + + println() + println("=== RESULT ===") + println(result) + println() + + // Wacht even zodat persistence kan schrijven + kotlinx.coroutines.delay(1000) + + // Check wat er is opgeslagen + val agentDir = File(persistenceDir, "checkpoints/" + agentId) + + println("=== PERSISTENCE CHECK ===") + if (agentDir.exists()) { + println("Snapshot directory exists: ${agentDir.absolutePath}") + agentDir.listFiles()?.forEach { + println(" - ${it.name} (${it.length()} bytes, modified: ${java.util.Date(it.lastModified())})") + } + } else { + println("WARNING: No snapshot directory found!") + } + println() + + println("=== TEST COMPLETED ===") + println("To continue this conversation, run:") + println(" ./gradlew run --args=\"$issueNumber 'your next message'\"") + + } catch (e: Exception) { + log.error("Error during test: ${e.message}", e) + e.printStackTrace() + } + } +} + +/** + * Usage examples: + * + * Run with default values: + * ./gradlew run + * + * Run with specific issue number: + * ./gradlew run --args="5" + * + * Run with issue number and message: + * ./gradlew run --args="5 'Tell me a joke'" + * + * Continue conversation (reuse same issue number): + * ./gradlew run --args="5 'Tell me another one'" + * + * This simple agent has: + * - No tools + * - No custom strategy + * - Only basic chat functionality + * - Persistence enabled + */ \ No newline at end of file diff --git a/src/main/kotlin/nl/trivion/softwarefabric/agents/DesignerAgent.kt b/src/main/kotlin/nl/trivion/softwarefabric/agents/DesignerAgent.kt index 8ca6321..4a3ec82 100644 --- a/src/main/kotlin/nl/trivion/softwarefabric/agents/DesignerAgent.kt +++ b/src/main/kotlin/nl/trivion/softwarefabric/agents/DesignerAgent.kt @@ -99,6 +99,14 @@ class DesignerAgent(private val giteaService: GiteaService) { (nodeSendToolResult forwardTo nodeExecuteTool) onToolCall { true } ) + + edge( + (nodeSendToolResult forwardTo nodeFinish) + onAssistantMessage { + log.info("content: ${it.content}") + it.content.startsWith("Question is send to") + } + ) } val systemPrompt = this::class.java.getResource("/prompts/designer-system-prompt.txt") diff --git a/src/main/kotlin/nl/trivion/softwarefabric/agents/SnapshotStrategy.kt b/src/main/kotlin/nl/trivion/softwarefabric/agents/SnapshotStrategy.kt new file mode 100644 index 0000000..c55858f --- /dev/null +++ b/src/main/kotlin/nl/trivion/softwarefabric/agents/SnapshotStrategy.kt @@ -0,0 +1,53 @@ +package nl.trivion.softwarefabric.agents + +import ai.koog.agents.core.dsl.builder.AIAgentNodeDelegate +import ai.koog.agents.core.dsl.builder.AIAgentSubgraphBuilderBase +import ai.koog.agents.core.dsl.builder.forwardTo +import ai.koog.agents.core.dsl.builder.strategy +import ai.koog.agents.snapshot.feature.withPersistence +import kotlinx.serialization.json.JsonPrimitive + +private fun AIAgentSubgraphBuilderBase<*, *>.simpleNode( + name: String? = null, + output: String, +): AIAgentNodeDelegate = node(name) { + return@node it + output +} + +private data class TeleportState(var teleported: Boolean = false) + +private fun AIAgentSubgraphBuilderBase<*, *>.teleportNode( + name: String? = null, + teleportState: TeleportState = TeleportState() +): AIAgentNodeDelegate = node(name) { + if (!teleportState.teleported) { + teleportState.teleported = true + withPersistence { + setExecutionPoint(it, "Node1", listOf(), JsonPrimitive("Teleported!!!")) + return@withPersistence "Teleported" + } + } else { + return@node "$it\nAlready teleported, passing by" + } +} + +object SnapshotStrategy { + private val teleportState = TeleportState() + + val strategy = strategy("test") { + val node1 by simpleNode( + "Node1", + output = "Node 1 output" + ) + val node2 by simpleNode( + "Node2", + output = "Node 2 output" + ) + val teleportNode by teleportNode(teleportState = teleportState) + + edge(nodeStart forwardTo node1) + edge(node1 forwardTo node2) + edge(node2 forwardTo teleportNode) + edge(teleportNode forwardTo nodeFinish) + } +} \ No newline at end of file diff --git a/src/main/kotlin/nl/trivion/softwarefabric/client/TestAgent.kt b/src/main/kotlin/nl/trivion/softwarefabric/client/TestAgent.kt new file mode 100644 index 0000000..621a7fa --- /dev/null +++ b/src/main/kotlin/nl/trivion/softwarefabric/client/TestAgent.kt @@ -0,0 +1,59 @@ +package nl.trivion.softwarefabric.client + +import ai.koog.a2a.client.A2AClient +import ai.koog.a2a.client.UrlAgentCardResolver +import ai.koog.a2a.consts.A2AConsts +import ai.koog.a2a.model.Message +import ai.koog.a2a.model.MessageSendConfiguration +import ai.koog.a2a.model.MessageSendParams +import ai.koog.a2a.model.Role +import ai.koog.a2a.model.TextPart +import ai.koog.a2a.transport.Request +import ai.koog.a2a.transport.client.jsonrpc.http.HttpJSONRPCClientTransport +import kotlin.uuid.ExperimentalUuidApi +import kotlin.uuid.Uuid + +@OptIn(ExperimentalUuidApi::class) +suspend fun main() { + val agentUrl = "http://localhost:8080/a2a" + + val cardResolver = UrlAgentCardResolver( + baseUrl = agentUrl, + path = A2AConsts.AGENT_CARD_WELL_KNOWN_PATH, + ) + + val transport = HttpJSONRPCClientTransport( + url = agentUrl, + ) + + val a2aClient = A2AClient( + transport = transport, + agentCardResolver = cardResolver + ) + +// Initialize client and fetch the card + println("connect") + a2aClient.connect() + println("connected") + + val message = Message( + role = Role.User, + parts = listOf( + TextPart("Hoe heet ik?") + ), + contextId = "blog-project-456", + messageId = Uuid.random().toString() + ) + + val request = Request( + data = MessageSendParams( + message = message, + configuration = MessageSendConfiguration( + blocking = false, // Get first response + historyLength = 5 // Include context + ) + ) + ) + val response = a2aClient.sendMessage(request) + println(response.data) +} \ No newline at end of file diff --git a/src/main/kotlin/nl/trivion/softwarefabric/controllers/A2AController2.kt b/src/main/kotlin/nl/trivion/softwarefabric/controllers/A2AController2.kt new file mode 100644 index 0000000..d7c6396 --- /dev/null +++ b/src/main/kotlin/nl/trivion/softwarefabric/controllers/A2AController2.kt @@ -0,0 +1,166 @@ +package nl.trivion.softwarefabric.controllers + +import ai.koog.a2a.consts.A2AConsts +import ai.koog.a2a.model.* +import ai.koog.a2a.server.A2AServer +import ai.koog.a2a.server.agent.AgentExecutor +import ai.koog.a2a.server.session.RequestContext +import ai.koog.a2a.server.session.SessionEventProcessor +import ai.koog.a2a.transport.server.jsonrpc.http.HttpJSONRPCServerTransport +import ai.koog.agents.a2a.core.A2AMessage +import ai.koog.agents.a2a.core.toA2AMessage +import ai.koog.agents.a2a.core.toKoogMessage +import ai.koog.agents.a2a.server.feature.A2AAgentServer +import ai.koog.agents.a2a.server.feature.nodeA2ARespondMessage +import ai.koog.agents.core.agent.AIAgent +import ai.koog.agents.core.agent.config.AIAgentConfig +import ai.koog.agents.core.dsl.builder.strategy +import ai.koog.agents.features.eventHandler.feature.handleEvents +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.executor.llms.all.simpleOllamaAIExecutor +import ai.koog.prompt.llm.LLMCapability +import ai.koog.prompt.llm.LLMProvider +import ai.koog.prompt.llm.LLModel +import io.ktor.server.cio.CIO +import nl.trivion.softwarefabric.storage.FileMessageStorage +import kotlin.uuid.ExperimentalUuidApi +import kotlin.uuid.Uuid + +class A2AController2 { + + @OptIn(ExperimentalUuidApi::class) + fun blogpostWritingStrategy() = strategy("blogpost-writer-strategy") { + val blogpostRequest by node { input -> + val userMessage = input.message.toKoogMessage().content + val contextId = input.message.contextId + val taskId = input.message.taskId + + llm.writeSession { + appendPrompt { user(userMessage) } + val llmResponse = requestLLM() + A2AMessage( + messageId = Uuid.random().toString(), + role = Role.Agent, + parts = listOf(TextPart(llmResponse.content)), + contextId = contextId, + taskId = taskId + ) + } + } + + val sendMessage by nodeA2ARespondMessage() + + nodeStart then blogpostRequest then sendMessage then nodeFinish + } + + private val customModel: LLModel = LLModel( + provider = LLMProvider.Ollama, + id = "gpt-oss:20b", + capabilities = listOf(LLMCapability.Temperature), + contextLength = 2048 + ) + + private val systemPrompt = """ + You are a helpful assistant. Keep your responses concise and friendly. + """.trimIndent() + + private val promptExecutor = simpleOllamaAIExecutor("http://localhost:11434") + + suspend fun createBlogpostWritingAgent( + requestContext: RequestContext, + eventProcessor: SessionEventProcessor + ): AIAgent { + requestContext.messageStorage.save(requestContext.params.message) + // De A2AServer heeft al het inkomende bericht opgeslagen via de messageStorageFactory + // Haal de volledige geschiedenis op (inclusief het huidige bericht) + val messageHistory = requestContext.messageStorage.getAll().map { it.toKoogMessage() } + + val agentConfig = AIAgentConfig( + prompt = Prompt.build("simple-chat") { + system(systemPrompt) + messages(messageHistory) + }, + model = customModel, + maxAgentIterations = 5 + ) + + return AIAgent( + promptExecutor = promptExecutor, + strategy = blogpostWritingStrategy(), + agentConfig = agentConfig + ) { + install(A2AAgentServer) { + this.context = requestContext + this.eventProcessor = eventProcessor + } + + handleEvents { + onAgentCompleted { ctx -> + val resultMessage = ctx.result as A2AMessage + // Sla het antwoord op via de messageStorage van de context + requestContext.messageStorage.save(resultMessage) + } + } + } + } + + inner class BlogpostAgentExecutor : AgentExecutor { + override suspend fun execute( + context: RequestContext, + eventProcessor: SessionEventProcessor + ) { + createBlogpostWritingAgent(context, eventProcessor) + .run(context.params) + } + } + + val agentCard = AgentCard( + name = "Blog Writer", + description = "AI agent that creates high-quality blog posts and articles", + url = "https://api.blog-writer.com/a2a/v1", + version = "1.0.0", + capabilities = AgentCapabilities(streaming = true), + defaultInputModes = listOf("text/plain"), + defaultOutputModes = listOf("text/markdown"), + skills = listOf( + AgentSkill( + id = "write-post", + name = "Blog Post Writing", + description = "Generate engaging blog posts on any topic", + tags = listOf("writing", "content", "blog"), + examples = listOf("Write a post about AI trends") + ) + ) + ) + + suspend fun start() { + // Maak een file-based message storage aan + val fileMessageStorage = FileMessageStorage( + storageDir = "./a2a-messages" + ) + + val a2aServer = A2AServer( + agentExecutor = BlogpostAgentExecutor(), + agentCard = agentCard, + messageStorage = fileMessageStorage + ) + + val transport = HttpJSONRPCServerTransport( + requestHandler = a2aServer + ) + + transport.start( + engineFactory = CIO, + port = 8080, + path = "/a2a", + wait = true, + agentCard = agentCard, + agentCardPath = A2AConsts.AGENT_CARD_WELL_KNOWN_PATH + ) + } +} + +suspend fun main() { + val controller = A2AController2() + controller.start() +} \ No newline at end of file diff --git a/src/main/kotlin/nl/trivion/softwarefabric/storage/FileMessageStorage.kt b/src/main/kotlin/nl/trivion/softwarefabric/storage/FileMessageStorage.kt new file mode 100644 index 0000000..357d10e --- /dev/null +++ b/src/main/kotlin/nl/trivion/softwarefabric/storage/FileMessageStorage.kt @@ -0,0 +1,172 @@ +package nl.trivion.softwarefabric.storage + +import ai.koog.a2a.model.Message +import ai.koog.a2a.server.messages.MessageStorage +import ai.koog.a2a.server.exceptions.MessageOperationException +import kotlinx.serialization.json.Json +import kotlinx.serialization.SerializationException +import java.io.File +import java.nio.file.Paths +import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock + +/** + * File-based implementation of MessageStorage for A2A server. + * Stores messages per contextId in separate JSON files. + * Thread-safe implementation using Mutex. + * + * Usage: + * ``` + * val messageStorage = FileMessageStorage(storageDir = "./a2a-messages") + * val a2aServer = A2AServer( + * agentExecutor = myExecutor, + * agentCard = myCard, + * messageStorage = messageStorage + * ) + * ``` + */ +class FileMessageStorage( + private val storageDir: String = "./a2a-messages" +) : MessageStorage { + + private val json = Json { + prettyPrint = true + ignoreUnknownKeys = true + } + + private val mutex = Mutex() + + init { + // Zorg ervoor dat de storage directory bestaat + File(storageDir).mkdirs() + } + + private fun getContextFile(contextId: String): File { + // Sanitize contextId voor bestandsnaam (verwijder ongeldige karakters) + val safeContextId = contextId.replace(Regex("[^a-zA-Z0-9._-]"), "_") + return Paths.get(storageDir, "$safeContextId.json").toFile() + } + + override suspend fun save(message: Message) { + val contextId = message.contextId ?: throw MessageOperationException( + "Cannot save message without contextId" + ) + + mutex.withLock { + try { + val file = getContextFile(contextId) + val messages = if (file.exists()) { + try { + val content = file.readText() + json.decodeFromString>(content) + } catch (e: SerializationException) { + throw MessageOperationException( + "Failed to parse existing messages for context $contextId", + e + ) + } + } else { + mutableListOf() + } + + messages.add(message) + + val jsonContent = json.encodeToString(messages) + file.writeText(jsonContent) + + } catch (e: Exception) { + if (e is MessageOperationException) throw e + throw MessageOperationException( + "Failed to save message for context $contextId", + e + ) + } + } + } + + override suspend fun getByContext(contextId: String): List { + return mutex.withLock { + try { + val file = getContextFile(contextId) + if (file.exists()) { + val content = file.readText() + json.decodeFromString>(content) + } else { + emptyList() + } + } catch (e: SerializationException) { + throw MessageOperationException( + "Failed to read messages for context $contextId", + e + ) + } catch (e: Exception) { + throw MessageOperationException( + "Failed to retrieve messages for context $contextId", + e + ) + } + } + } + + override suspend fun deleteByContext(contextId: String) { + mutex.withLock { + try { + val file = getContextFile(contextId) + if (file.exists() && !file.delete()) { + throw MessageOperationException( + "Failed to delete messages for context $contextId" + ) + } + } catch (e: Exception) { + if (e is MessageOperationException) throw e + throw MessageOperationException( + "Failed to delete messages for context $contextId", + e + ) + } + } + } + + override suspend fun replaceByContext(contextId: String, messages: List) { + mutex.withLock { + try { + // Valideer dat alle messages de juiste contextId hebben + if (messages.any { it.contextId != contextId }) { + throw MessageOperationException( + "All messages must have contextId '$contextId'" + ) + } + + val file = getContextFile(contextId) + val jsonContent = json.encodeToString(messages) + file.writeText(jsonContent) + + } catch (e: Exception) { + if (e is MessageOperationException) throw e + throw MessageOperationException( + "Failed to replace messages for context $contextId", + e + ) + } + } + } + + /** + * Utility: Clear all stored messages across all contexts + */ + suspend fun clearAll() { + mutex.withLock { + File(storageDir).listFiles()?.forEach { it.delete() } + } + } + + /** + * Utility: Get list of all context IDs that have stored messages + */ + fun getAllContextIds(): List { + return File(storageDir).listFiles() + ?.filter { it.extension == "json" } + ?.map { it.nameWithoutExtension } + ?: emptyList() + } +} \ No newline at end of file