updates

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 {