Save the downstream service state after every interaction

This provides a basic implementation for #159

Author: orpiske

api/src/main/java/ai/wanaku/api/types/management/State.java

@@ -0,0 +1,4 @@
+package ai.wanaku.api.types.management;
+
+public record State(String service, boolean healthy, String message) {
+}

cli/src/main/java/ai/wanaku/cli/main/commands/targets/AbstractTargetState.java

@@ -0,0 +1,25 @@
+package ai.wanaku.cli.main.commands.targets;
+
+import java.net.URI;
+
+import ai.wanaku.cli.main.commands.BaseCommand;
+import ai.wanaku.cli.main.services.LinkService;
+import io.quarkus.rest.client.reactive.QuarkusRestClientBuilder;
+import picocli.CommandLine;
+
+@CommandLine.Command(name = "state",
+        description = "Get the state of the targeted services")
+public abstract class AbstractTargetState extends BaseCommand {
+    @CommandLine.Option(names = {"--host"}, description = "The API host", defaultValue = "http://localhost:8080",
+            arity = "0..1")
+    protected String host;
+
+    protected LinkService linkService;
+
+    protected void initService() {
+        linkService = QuarkusRestClientBuilder.newBuilder()
+                .baseUri(URI.create(host))
+                .build(LinkService.class);
+    }
+
+}

cli/src/main/java/ai/wanaku/cli/main/commands/targets/resources/Resources.java

@@ -4,7 +4,7 @@
 import picocli.CommandLine;
 
 @CommandLine.Command(name = "resources",
-        description = "Manage targets", subcommands = { ResourcesLinkedList.class, ResourcesConfigure.class })
+        description = "Manage targets", subcommands = { ResourcesLinkedList.class, ResourcesConfigure.class, ResourcesState.class })
 public class Resources extends BaseCommand {
     @Override
     public void run() {

cli/src/main/java/ai/wanaku/cli/main/commands/targets/resources/ResourcesState.java

@@ -0,0 +1,22 @@
+package ai.wanaku.cli.main.commands.targets.resources;
+
+import java.util.List;
+import java.util.Map;
+
+import ai.wanaku.api.types.WanakuResponse;
+import ai.wanaku.api.types.management.State;
+import ai.wanaku.cli.main.commands.targets.AbstractTargetsList;
+import ai.wanaku.cli.main.support.PrettyPrinter;
+import picocli.CommandLine;
+
+@CommandLine.Command(name = "state",
+        description = "List service states")
+public class ResourcesState extends AbstractTargetsList {
+    @Override
+    public void run() {
+        initService();
+
+        WanakuResponse<Map<String, List<State>>> list = linkService.resourcesState();
+        PrettyPrinter.printStates(list.data());
+    }
+}

cli/src/main/java/ai/wanaku/cli/main/commands/targets/tools/Tools.java

@@ -4,7 +4,7 @@
 import picocli.CommandLine;
 
 @CommandLine.Command(name = "tools",
-        description = "Manage targets", subcommands = { ToolsLinkedList.class, ToolsConfigure.class })
+        description = "Manage targets", subcommands = { ToolsLinkedList.class, ToolsConfigure.class, ToolsState.class })
 public class Tools extends BaseCommand {
     @Override
     public void run() {

cli/src/main/java/ai/wanaku/cli/main/commands/targets/tools/ToolsState.java

@@ -0,0 +1,21 @@
+package ai.wanaku.cli.main.commands.targets.tools;
+
+import java.util.List;
+import java.util.Map;
+
+import ai.wanaku.api.types.management.State;
+import ai.wanaku.cli.main.commands.targets.AbstractTargetsList;
+import ai.wanaku.cli.main.support.PrettyPrinter;
+import picocli.CommandLine;
+
+@CommandLine.Command(name = "state",
+        description = "List services states")
+public class ToolsState extends AbstractTargetsList {
+    @Override
+    public void run() {
+        initService();
+
+        Map<String, List<State>> states = linkService.toolsState().data();
+        PrettyPrinter.printStates(states);
+    }
+}

cli/src/main/java/ai/wanaku/cli/main/services/LinkService.java

@@ -1,8 +1,10 @@
 package ai.wanaku.cli.main.services;
 
+import java.util.List;
 import java.util.Map;
 
 import ai.wanaku.api.types.WanakuResponse;
+import ai.wanaku.api.types.management.State;
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.PUT;
@@ -22,6 +24,11 @@ public interface LinkService {
     @Consumes(MediaType.TEXT_PLAIN)
     WanakuResponse<Map<String, Service>> toolsList();
 
+    @Path("/tools/state")
+    @GET
+    @Consumes(MediaType.TEXT_PLAIN)
+    WanakuResponse<Map<String, List<State>>> toolsState();
+
     @Path("/tools/configure/{service}")
     @PUT
     @Consumes(MediaType.TEXT_PLAIN)
@@ -36,4 +43,9 @@ public interface LinkService {
     @PUT
     @Consumes(MediaType.TEXT_PLAIN)
     Response resourcesConfigure(@RestPath("service") String service, @QueryParam("option") String option, @QueryParam("value") String value);
+
+    @Path("/resources/state")
+    @GET
+    @Consumes(MediaType.TEXT_PLAIN)
+    WanakuResponse<Map<String, List<State>>> resourcesState();
 }

cli/src/main/java/ai/wanaku/cli/main/support/PrettyPrinter.java

@@ -8,6 +8,7 @@
 import ai.wanaku.api.types.ToolReference;
 import ai.wanaku.api.types.management.Configuration;
 import ai.wanaku.api.types.management.Service;
+import ai.wanaku.api.types.management.State;
 
 public class PrettyPrinter {
 
@@ -68,4 +69,26 @@ public static void printTargets(final Map<String, Service> map) {
         }
     }
 
+
+    /**
+     * Prints a map of entries
+     * @param states the map of states
+     */
+    public static void printStates(final Map<String, List<State>> states) {
+        System.out.printf("%-20s    %-10s    %-60s%n",
+                "Service", "Healthy", "Message");
+
+        for (var entry : states.entrySet()) {
+            for (var state : entry.getValue()) {
+                printParseableState(entry.getKey(), state.healthy(), state.message());
+            }
+        }
+    }
+
+    private static void printParseableState(String service, boolean healthy, String message) {
+        System.out.printf("%-15s => %-15s => %-30s    %n",
+                service, Boolean.valueOf(healthy), message);
+    }
+
+
 }

core/core-mcp/src/main/java/ai/wanaku/core/mcp/providers/ServiceRegistry.java

@@ -1,6 +1,9 @@
 package ai.wanaku.core.mcp.providers;
 
 import ai.wanaku.api.types.management.Service;
+import ai.wanaku.api.types.management.State;
+
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -19,7 +22,7 @@ public interface ServiceRegistry {
      * De-register a service from the registry
      * @param service the service name
      */
-    void deregister(String service);
+    void deregister(String service, ServiceType serviceType);
 
     /**
      * Gets a registered service by name
@@ -28,6 +31,22 @@ public interface ServiceRegistry {
      */
     Service getService(String service);
 
+    /**
+     * Saves the current state of the service
+     * @param service the service to save the state
+     * @param healthy whether it is healthy (true for healthy, false otherwise)
+     * @param message Optional state message (ignored if healthy)
+     */
+    void saveState(String service, boolean healthy, String message);
+
+    /**
+     * Gets the state of the given service
+     * @param service the service name
+     * @param count the number of states to get
+     * @return the last count states of the given service
+     */
+    List<State> getState(String service, int count);
+
     /**
      * Get a map of all registered services and their configurations
      * @param serviceType the type of service to get

core/core-service-discovery/src/main/java/ai/wanaku/core/service/discovery/ReservedKeys.java

@@ -2,6 +2,8 @@
 
 import java.util.Set;
 
+import ai.wanaku.core.mcp.providers.ServiceType;
+
 /**
  * Reserved keys
  */
@@ -16,6 +18,25 @@ class ReservedKeys {
      */
     public static final String WANAKU_TARGET_TYPE = "wanaku-target-type";
 
+    /**
+     * This key stores the set of Wanaku resource services
+     */
+    public static final String WANAKU_SERVICES_RESOURCES = "wanaku-services-resources";
+
+    /**
+     * This key stores the set of Wanaku tools services
+     */
+    public static final String WANAKU_SERVICES_TOOLS = "wanaku-services-tools";
+
 
     public static final Set<String> ALL_KEYS = Set.of(WANAKU_TARGET_ADDRESS, WANAKU_TARGET_TYPE);
+
+    public static String getServiceKey(ServiceType serviceType) {
+        switch (serviceType) {
+            case TOOL_INVOKER: return ReservedKeys.WANAKU_TARGET_ADDRESS;
+            case RESOURCE_PROVIDER: return ReservedKeys.WANAKU_SERVICES_RESOURCES;
+        }
+
+        throw new IllegalArgumentException("Unknown service type: " + serviceType);
+    }
 }

core/core-service-discovery/src/main/java/ai/wanaku/core/service/discovery/ValkeyRegistry.java

@@ -1,5 +1,9 @@
 package ai.wanaku.core.service.discovery;
 
+import ai.wanaku.api.types.management.State;
+import io.valkey.StreamEntryID;
+import io.valkey.params.XAddParams;
+import io.valkey.resps.StreamEntry;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.enterprise.event.Observes;
 import jakarta.inject.Inject;
@@ -13,7 +17,11 @@
 import io.quarkus.runtime.ShutdownEvent;
 import io.valkey.Jedis;
 import io.valkey.JedisPool;
+
+import java.time.Instant;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import org.jboss.logging.Logger;
@@ -46,6 +54,10 @@ public class ValkeyRegistry implements ServiceRegistry {
     @Override
     public void register(ServiceTarget serviceTarget, Map<String, String> configurations) {
         try (io.valkey.Jedis jedis = jedisPool.getResource()) {
+            // Register the service on the specific set
+            String serviceKey = ReservedKeys.getServiceKey(serviceTarget.getServiceType());
+            jedis.sadd(serviceKey, serviceTarget.getService());
+
             jedis.hset(serviceTarget.getService(), ReservedKeys.WANAKU_TARGET_ADDRESS, serviceTarget.toAddress());
             jedis.hset(serviceTarget.getService(), ReservedKeys.WANAKU_TARGET_TYPE, serviceTarget.getServiceType().asValue());
 
@@ -64,17 +76,70 @@ public void register(ServiceTarget serviceTarget, Map<String, String> configurat
      * Deregisters a service with the given name.
      *
      * @param service The name of the service to deregister.
+     * @param serviceType the type of service to deregister
      */
     @Override
-    public void deregister(String service) {
+    public void deregister(String service, ServiceType serviceType) {
         try (io.valkey.Jedis jedis = jedisPool.getResource()) {
-            jedis.del(service);
+            String serviceKey = ReservedKeys.getServiceKey(serviceType);
+            jedis.srem(serviceKey, service);
+
             LOG.infof("Service %s registered", service);
         } catch (Exception e) {
             LOG.errorf(e, "Failed to register service %s: %s", service, e.getMessage());
         }
     }
 
+
+    @Override
+    public void saveState(String service, boolean healthy, String message) {
+        try (io.valkey.Jedis jedis = jedisPool.getResource()) {
+            Map<String, String> state = Map.of("service", service, "healthy",
+                    Boolean.toString(healthy), "message", (healthy ? "healthy" : message));
+
+            jedis.xadd(stateKey(service), state, XAddParams.xAddParams());
+        } catch (Exception e) {
+            LOG.errorf(e, "Failed to save state for %s: %s", service, e.getMessage());
+        }
+    }
+
+    @Override
+    public List<State> getState(String service, int count) {
+        try (io.valkey.Jedis jedis = jedisPool.getResource()) {
+
+            String stateKey = stateKey(service);
+            Instant now = Instant.now();
+            long endEpoch = now.toEpochMilli();
+            long startEpoch = now.minusSeconds(60).toEpochMilli();
+
+            List<StreamEntry> streamEntries = jedis.xrange(stateKey, new StreamEntryID(startEpoch), new StreamEntryID(endEpoch));
+
+            List<State> states = new ArrayList<>(streamEntries.size());
+
+            for (StreamEntry streamEntry : streamEntries) {
+                LOG.debugf("Entry %s", streamEntry);
+
+                Map<String, String> fields = streamEntry.getFields();
+                String serviceName = fields.get("service");
+                String message = fields.get("message");
+                String healthy = fields.get("healthy");
+
+                State state = new State(serviceName, Boolean.parseBoolean(healthy), message);
+                states.add(state);
+            }
+
+            return states;
+        } catch (Exception e) {
+            LOG.errorf(e, "Failed to get state for %s: %s", service, e.getMessage());
+        }
+
+        return List.of();
+    }
+
+    private static String stateKey(String service) {
+        return "state:" + service;
+    }
+
     /**
      * Retrieves a service with the given name.
      *
@@ -98,14 +163,13 @@ public Service getService(String service) {
     public Map<String, Service> getEntries(ServiceType serviceType) {
         Map<String, Service> entries = new HashMap<>();
         try (io.valkey.Jedis jedis = jedisPool.getResource()) {
-            Set<String> keys = jedis.keys("*");
-            for (String key : keys) {
-                String sType = jedis.hget(key, ReservedKeys.WANAKU_TARGET_TYPE);
-                if (serviceType.asValue().equals(sType)) {
-                    Service service = newService(jedis, key);
-
-                    entries.put(key, service);
-                }
+            String serviceKey = ReservedKeys.getServiceKey(serviceType);
+            Set<String> services = jedis.smembers(serviceKey);
+
+            for (String key : services) {
+                Service service = newService(jedis, key);
+
+                entries.put(key, service);
             }
 
         } catch (Exception e) {