updates

server/internal/api/handlers.go

@@ -122,18 +122,32 @@ func deleteServer(c *gin.Context) {
 }
 
 func generateKey(c *gin.Context) {
-	// The agent triggers key generation itself; this endpoint signals
-	// the intent by returning the server so the caller knows to wait
-	// for the agent to upload via gRPC UploadGeneratedKey.
 	id := c.Param("id")
+
+	var body struct {
+		Label string `json:"label"`
+	}
+	_ = c.ShouldBindJSON(&body)
+	if body.Label == "" {
+		body.Label = "generated"
+	}
+
 	s, err := services.GetServer(id)
 	if err != nil {
 		c.JSON(http.StatusNotFound, gin.H{"error": "server not found"})
 		return
 	}
-	c.JSON(http.StatusOK, gin.H{
-		"message":   "agent will generate and upload key on next poll",
-		"server_id": s.ServerID,
+
+	cmdID, err := services.DispatchGenerateKey(s.ServerID, body.Label)
+	if err != nil {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(http.StatusAccepted, gin.H{
+		"message":    "key generation command sent to agent",
+		"command_id": cmdID,
+		"server_id":  s.ServerID,
 	})
 }
 

server/internal/grpc/pb/keymanager.pb.go

@@ -45,12 +45,87 @@ type UploadKeyResponse struct {
 	KeyId string `json:"key_id"`
 }
 
+// CommandStream message types
+
+type ServerCommand struct {
+	CommandId   string          `json:"command_id"`
+	GenerateKey *GenerateKeyCmd `json:"generate_key,omitempty"`
+}
+
+type GenerateKeyCmd struct {
+	Label string `json:"label"`
+}
+
+type AgentMessage struct {
+	ServerId   string         `json:"server_id"`
+	AgentToken string         `json:"agent_token"`
+	Ready      *AgentReady    `json:"ready,omitempty"`
+	Result     *CommandResult `json:"result,omitempty"`
+}
+
+type AgentReady struct{}
+
+type CommandResult struct {
+	CommandId string `json:"command_id"`
+	Success   bool   `json:"success"`
+	Message   string `json:"message"`
+}
+
+// CommandStream server-side interface
+
+type KeyManager_CommandStreamServer interface {
+	Send(*ServerCommand) error
+	Recv() (*AgentMessage, error)
+	grpc.ServerStream
+}
+
+type keyManagerCommandStreamServer struct {
+	grpc.ServerStream
+}
+
+func (s *keyManagerCommandStreamServer) Send(m *ServerCommand) error {
+	return s.ServerStream.SendMsg(m)
+}
+
+func (s *keyManagerCommandStreamServer) Recv() (*AgentMessage, error) {
+	m := new(AgentMessage)
+	if err := s.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// CommandStream client-side interface
+
+type KeyManager_CommandStreamClient interface {
+	Send(*AgentMessage) error
+	Recv() (*ServerCommand, error)
+	grpc.ClientStream
+}
+
+type keyManagerCommandStreamClient struct {
+	grpc.ClientStream
+}
+
+func (c *keyManagerCommandStreamClient) Send(m *AgentMessage) error {
+	return c.ClientStream.SendMsg(m)
+}
+
+func (c *keyManagerCommandStreamClient) Recv() (*ServerCommand, error) {
+	m := new(ServerCommand)
+	if err := c.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 // Server interface
 
 type KeyManagerServer interface {
 	Register(context.Context, *RegisterRequest) (*RegisterResponse, error)
 	SyncKeys(context.Context, *SyncRequest) (*SyncResponse, error)
 	UploadGeneratedKey(context.Context, *UploadKeyRequest) (*UploadKeyResponse, error)
+	CommandStream(KeyManager_CommandStreamServer) error
 }
 
 type UnimplementedKeyManagerServer struct{}
@@ -67,12 +142,17 @@ func (UnimplementedKeyManagerServer) UploadGeneratedKey(context.Context, *Upload
 	return nil, status.Errorf(codes.Unimplemented, "method UploadGeneratedKey not implemented")
 }
 
+func (UnimplementedKeyManagerServer) CommandStream(KeyManager_CommandStreamServer) error {
+	return status.Errorf(codes.Unimplemented, "method CommandStream not implemented")
+}
+
 // Client interface
 
 type KeyManagerClient interface {
 	Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error)
 	SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error)
 	UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error)
+	CommandStream(ctx context.Context, opts ...grpc.CallOption) (KeyManager_CommandStreamClient, error)
 }
 
 type keyManagerClient struct {
@@ -107,6 +187,14 @@ func (c *keyManagerClient) UploadGeneratedKey(ctx context.Context, in *UploadKey
 	return out, nil
 }
 
+func (c *keyManagerClient) CommandStream(ctx context.Context, opts ...grpc.CallOption) (KeyManager_CommandStreamClient, error) {
+	stream, err := c.cc.NewStream(ctx, &KeyManager_ServiceDesc.Streams[0], "/keymanager.v1.KeyManager/CommandStream", opts...)
+	if err != nil {
+		return nil, err
+	}
+	return &keyManagerCommandStreamClient{stream}, nil
+}
+
 // Server registration
 
 func RegisterKeyManagerServer(s grpc.ServiceRegistrar, srv KeyManagerServer) {
@@ -121,7 +209,14 @@ var KeyManager_ServiceDesc = grpc.ServiceDesc{
 		{MethodName: "SyncKeys", Handler: _KeyManager_SyncKeys_Handler},
 		{MethodName: "UploadGeneratedKey", Handler: _KeyManager_UploadGeneratedKey_Handler},
 	},
-	Streams:  []grpc.StreamDesc{},
+	Streams: []grpc.StreamDesc{
+		{
+			StreamName:    "CommandStream",
+			Handler:       _KeyManager_CommandStream_Handler,
+			ServerStreams: true,
+			ClientStreams: true,
+		},
+	},
 	Metadata: "keymanager/v1/keymanager.proto",
 }
 
@@ -169,3 +264,7 @@ func _KeyManager_UploadGeneratedKey_Handler(srv interface{}, ctx context.Context
 	}
 	return interceptor(ctx, in, info, handler)
 }
+
+func _KeyManager_CommandStream_Handler(srv interface{}, stream grpc.ServerStream) error {
+	return srv.(KeyManagerServer).CommandStream(&keyManagerCommandStreamServer{stream})
+}

server/internal/grpc/server.go

@@ -67,6 +67,59 @@ func (s *keyManagerServer) UploadGeneratedKey(ctx context.Context, req *pb.Uploa
 	return &pb.UploadKeyResponse{KeyId: key.KeyID}, nil
 }
 
+func (s *keyManagerServer) CommandStream(stream pb.KeyManager_CommandStreamServer) error {
+	// First message authenticates the agent and signals readiness.
+	msg, err := stream.Recv()
+	if err != nil {
+		return status.Errorf(codes.InvalidArgument, "expected initial auth message: %v", err)
+	}
+
+	srv, err := services.ValidateAgentToken(msg.ServerId, msg.AgentToken)
+	if err != nil {
+		return status.Errorf(codes.Unauthenticated, "invalid agent token")
+	}
+
+	if err := services.UpdateServerLastSeen(srv.ServerID); err != nil {
+		log.Printf("update last seen %s: %v", srv.ServerID, err)
+	}
+
+	ch := services.Dispatcher.Connect(srv.ServerID)
+	defer services.Dispatcher.Disconnect(srv.ServerID)
+
+	log.Printf("agent %s connected command stream", srv.ServerID)
+	defer log.Printf("agent %s disconnected command stream", srv.ServerID)
+
+	// Drain inbound results in the background so client Send calls never block.
+	// UploadGeneratedKey handles the real storage; these are just confirmation logs.
+	go func() {
+		for {
+			m, err := stream.Recv()
+			if err != nil {
+				return
+			}
+			if m.Result != nil {
+				r := m.Result
+				log.Printf("agent %s cmd %s: success=%v %s", srv.ServerID, r.CommandId, r.Success, r.Message)
+			}
+		}
+	}()
+
+	ctx := stream.Context()
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		case cmd, ok := <-ch:
+			if !ok {
+				return nil
+			}
+			if err := stream.Send(cmd); err != nil {
+				return err
+			}
+		}
+	}
+}
+
 func StartGRPC(port int) error {
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {

server/internal/services/dispatch.go

@@ -0,0 +1,76 @@
+package services
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/google/uuid"
+	"github.com/mrhid6/keymanager/server/internal/grpc/pb"
+)
+
+type commandDispatcher struct {
+	mu       sync.RWMutex
+	channels map[string]chan *pb.ServerCommand
+}
+
+// Dispatcher is the singleton command dispatcher used by both the gRPC server
+// and the REST API to push commands to connected agents.
+var Dispatcher = &commandDispatcher{
+	channels: make(map[string]chan *pb.ServerCommand),
+}
+
+// Connect registers an agent's command channel. Returns the channel to drain.
+func (d *commandDispatcher) Connect(serverID string) chan *pb.ServerCommand {
+	ch := make(chan *pb.ServerCommand, 16)
+	d.mu.Lock()
+	d.channels[serverID] = ch
+	d.mu.Unlock()
+	return ch
+}
+
+// Disconnect removes the agent's channel on stream close.
+func (d *commandDispatcher) Disconnect(serverID string) {
+	d.mu.Lock()
+	delete(d.channels, serverID)
+	d.mu.Unlock()
+}
+
+// IsConnected reports whether an agent is currently holding a CommandStream.
+func (d *commandDispatcher) IsConnected(serverID string) bool {
+	d.mu.RLock()
+	_, ok := d.channels[serverID]
+	d.mu.RUnlock()
+	return ok
+}
+
+func (d *commandDispatcher) dispatch(serverID string, cmd *pb.ServerCommand) error {
+	d.mu.RLock()
+	ch, ok := d.channels[serverID]
+	d.mu.RUnlock()
+	if !ok {
+		return fmt.Errorf("agent for server %s is not connected", serverID)
+	}
+	select {
+	case ch <- cmd:
+		return nil
+	default:
+		return fmt.Errorf("command queue full for server %s", serverID)
+	}
+}
+
+// DispatchGenerateKey sends a generate-key command to the named server's agent.
+// Returns the command ID that can be used to correlate the agent's result.
+func DispatchGenerateKey(serverID, label string) (string, error) {
+	if !Dispatcher.IsConnected(serverID) {
+		return "", fmt.Errorf("agent is not connected to the command stream")
+	}
+	cmdID := uuid.New().String()
+	cmd := &pb.ServerCommand{
+		CommandId:   cmdID,
+		GenerateKey: &pb.GenerateKeyCmd{Label: label},
+	}
+	if err := Dispatcher.dispatch(serverID, cmd); err != nil {
+		return "", err
+	}
+	return cmdID, nil
+}