From 333d087979bb23ab12fd4ffd7347aecc1bdf9774 Mon Sep 17 00:00:00 2001 From: slhafzjw Date: Fri, 27 Feb 2026 00:28:14 +0800 Subject: [PATCH] refactor(framework): register and execute ordered `@Shutdown` hooks in `AgentContext` via JVM shutdown hook --- .../api/agent/factory/context/AgentContext.kt | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt b/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt index 0fd9b2ac..9651135b 100644 --- a/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt +++ b/Partner-Framework/src/main/java/work/slhaf/partner/api/agent/factory/context/AgentContext.kt @@ -2,8 +2,12 @@ package work.slhaf.partner.api.agent.factory.context import com.alibaba.fastjson2.JSONArray import com.alibaba.fastjson2.JSONObject +import org.slf4j.LoggerFactory +import work.slhaf.partner.api.agent.factory.capability.annotation.CapabilityCore import work.slhaf.partner.api.agent.factory.component.abstracts.AbstractAgentModule import work.slhaf.partner.api.agent.factory.component.annotation.AgentComponent +import work.slhaf.partner.api.agent.runtime.exception.AgentRunningFailedException +import work.slhaf.partner.api.agent.runtime.interaction.flow.entity.RunningFlowContext import java.lang.reflect.Method import java.time.ZonedDateTime @@ -31,6 +35,12 @@ object AgentContext { val metadata: Map get() = _metadata + private val shutdownHooks = mutableMapOf>() + + init { + installShutdownHook() + } + fun addModule(name: String, module: ModuleContextData) { _modules[name] = module } @@ -59,6 +69,104 @@ object AgentContext { _metadata[name] = content } + fun addShutdownHook(method: Method, order: Int): Boolean { + if (!method.isAnnotationPresent(Shutdown::class.java)) { + return false + } + val clazz = method.declaringClass + val type = if (AbstractAgentModule.Running::class.java.isAssignableFrom(clazz)) { + ShutdownHookDesc.Type.RUNNING + } else if (AbstractAgentModule.Sub::class.java.isAssignableFrom(clazz)) { + ShutdownHookDesc.Type.SUB + } else if (AbstractAgentModule.Standalone::class.java.isAssignableFrom(clazz)) { + ShutdownHookDesc.Type.STANDALONE + } else if (clazz.isAnnotationPresent(AgentComponent::class.java)) { + ShutdownHookDesc.Type.ADDITIONAL + } else if (clazz.isAnnotationPresent(CapabilityCore::class.java)) { + ShutdownHookDesc.Type.CAPABILITY + } else { + return false + } + shutdownHooks.computeIfAbsent(type) { mutableListOf() }.add(ShutdownHookDesc(clazz, order, method, type)) + return true + } + + private fun installShutdownHook() { + + fun getModuleInstance(clazz: Class<*>, instances: Instances): Any? { + return if (AbstractAgentModule.Running::class.java.isAssignableFrom(clazz)) { + instances.running[clazz] + } else if (AbstractAgentModule.Sub::class.java.isAssignableFrom(clazz)) { + instances.sub[clazz] + } else if (AbstractAgentModule.Standalone::class.java.isAssignableFrom(clazz)) { + instances.standalone[clazz] + } else { + null + } + } + + fun getInstanceOf(clazz: Class<*>, type: ShutdownHookDesc.Type, instances: Instances): Any { + val instance = when (type) { + ShutdownHookDesc.Type.RUNNING -> getModuleInstance(clazz, instances) + ShutdownHookDesc.Type.STANDALONE -> getModuleInstance(clazz, instances) + ShutdownHookDesc.Type.SUB -> getModuleInstance(clazz, instances) + ShutdownHookDesc.Type.ADDITIONAL -> instances.additional[clazz] + ShutdownHookDesc.Type.CAPABILITY -> instances.capability[clazz] + } + if (instance == null) { + throw AgentRunningFailedException("Instance of type $type not found") + } + return instance + } + + fun trigger(hooks: List, instances: Instances) { + val log = LoggerFactory.getLogger(AgentContext::class.java) + hooks.sortedBy { it.order } + .forEach { + try { + it.method.invoke(getInstanceOf(it.clazz, it.type, instances)) + } catch (e: Exception) { + log.error("Failed to invoke shutdown hook ${it.clazz.simpleName}#${it.method.name}", e) + } + } + } + + Runtime.getRuntime().addShutdownHook(Thread { + val instances = computeInstances() + shutdownHooks[ShutdownHookDesc.Type.RUNNING]?.let { trigger(it, instances) } + shutdownHooks[ShutdownHookDesc.Type.ADDITIONAL]?.let { trigger(it, instances) } + shutdownHooks[ShutdownHookDesc.Type.STANDALONE]?.let { trigger(it, instances) } + shutdownHooks[ShutdownHookDesc.Type.SUB]?.let { trigger(it, instances) } + shutdownHooks[ShutdownHookDesc.Type.CAPABILITY]?.let { trigger(it, instances) } + }) + } + + private fun computeInstances(): Instances { + val instances = Instances() + modules.values.forEach { + when (it) { + is ModuleContextData.Running<*> -> instances.running[it.clazz] = it.instance + is ModuleContextData.Standalone<*> -> instances.standalone[it.clazz] = it.instance + is ModuleContextData.Sub<*> -> instances.sub[it.clazz] = it.instance + } + + } + instances.additional.putAll(additionalComponents) + capabilities.values.forEach { + instances.capability.putAll(it.cores) + } + return instances + } + + private class Instances( + val running: MutableMap>, Any> = mutableMapOf(), + val standalone: MutableMap, Any> = mutableMapOf(), + val sub: MutableMap>, Any> = mutableMapOf(), + val additional: MutableMap, Any> = mutableMapOf(), + val capability: MutableMap, Any> = mutableMapOf() + ) + + data class MetaDataContent( val clazz: Class<*>, val value: String @@ -126,3 +234,18 @@ sealed class ModuleContextData { annotation class Shutdown( val order: Int = 0, ) + +data class ShutdownHookDesc( + val clazz: Class<*>, + val order: Int, + val method: Method, + val type: Type +) { + enum class Type { + RUNNING, + ADDITIONAL, + STANDALONE, + SUB, + CAPABILITY + } +}